taichi.lang
1# type: ignore 2 3from taichi.lang import impl, simt 4from taichi.lang._ndarray import * 5from taichi.lang._ndrange import ndrange 6from taichi.lang._texture import Texture 7from taichi.lang.argpack import * 8from taichi.lang.exception import * 9from taichi.lang.field import * 10from taichi.lang.impl import * 11from taichi.lang.kernel_impl import * 12from taichi.lang.matrix import * 13from taichi.lang.mesh import * 14from taichi.lang.misc import * # pylint: disable=W0622 15from taichi.lang.ops import * # pylint: disable=W0622 16from taichi.lang.runtime_ops import * 17from taichi.lang.snode import * 18from taichi.lang.source_builder import * 19from taichi.lang.struct import * 20from taichi.types.enums import DeviceCapability, Format, Layout 21 22__all__ = [ 23 s 24 for s in dir() 25 if not s.startswith("_") 26 and s 27 not in [ 28 "any_array", 29 "ast", 30 "common_ops", 31 "enums", 32 "exception", 33 "expr", 34 "impl", 35 "inspect", 36 "kernel_arguments", 37 "kernel_impl", 38 "matrix", 39 "mesh", 40 "misc", 41 "ops", 42 "platform", 43 "runtime_ops", 44 "shell", 45 "snode", 46 "source_builder", 47 "struct", 48 "util", 49 ] 50]
26class ArgPack: 27 """ The `ArgPack` Type Class. 28 29 The `ArgPack` operates as a dictionary-like data pack, storing members as (key, value) pairs. Members stored can 30 range from scalars and matrices to other dictionary-like structures. Distinguished from structs, `ArgPack` can 31 accommodate buffer types such as `NdarrayType` and `TextureType` from Taichi. However, unlike `ti.Struct` which 32 serves as a data container, `ArgPack` functions as a reference container. It's important to note that `ArgPack` 33 cannot be nested within other types except for another `ArgPack`, and can only be utilized as kernel parameters. 34 35 Args: 36 annotations (Dict[str, Union[Dict, Matrix, Struct]]): \ 37 The keys and types for `ArgPack` members. 38 dtype (ArgPackType): \ 39 The ArgPackType class of this ArgPack object. 40 entries (Dict[str, Union[Dict, Matrix, Struct]]): \ 41 The keys and corresponding values for `ArgPack` members. 42 43 Returns: 44 An instance of this `ArgPack`. 45 46 Example:: 47 48 >>> vec3 = ti.types.vector(3, ti.f32) 49 >>> pack_type = ti.ArgPackType(v=vec3, t=ti.f32) 50 >>> a = pack_type(v=vec3([0, 0, 0]), t=1.0) 51 >>> print(a.items) 52 dict_items([('v', [0. 0. 0.]), ('t', 1.0)]) 53 """ 54 55 _instance_count = 0 56 57 def __init__(self, annotations, dtype, *args, **kwargs): 58 # converts dicts to argument packs 59 if len(args) == 1 and kwargs == {} and isinstance(args[0], dict): 60 self.__entries = args[0] 61 elif len(args) == 0: 62 self.__entries = kwargs 63 else: 64 raise TaichiSyntaxError( 65 "Custom argument packs need to be initialized using either dictionary or keyword arguments" 66 ) 67 if annotations.keys() != self.__entries.keys(): 68 raise TaichiSyntaxError("ArgPack annotations keys not equals to entries keys.") 69 self.__annotations = annotations 70 for k, v in self.__entries.items(): 71 self.__entries[k] = v if in_python_scope() else impl.expr_init(v) 72 self._register_members() 73 self.__dtype = dtype 74 self.__argpack = impl.get_runtime().prog.create_argpack(self.__dtype) 75 for i, (k, v) in enumerate(self.__entries.items()): 76 self._write_to_device(self.__annotations[k], type(v), v, self._calc_element_true_index(i)) 77 78 def __del__(self): 79 if impl is not None and impl.get_runtime() is not None and impl.get_runtime().prog is not None: 80 impl.get_runtime().prog.delete_argpack(self.__argpack) 81 82 @property 83 def keys(self): 84 """Returns the list of member names in string format. 85 86 Example:: 87 88 >>> vec3 = ti.types.vector(3, ti.f32) 89 >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32) 90 >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0) 91 >>> sphere.keys 92 ['center', 'radius'] 93 """ 94 return list(self.__entries.keys()) 95 96 @property 97 def _members(self): 98 return list(self.__entries.values()) 99 100 @property 101 def _annotations(self): 102 return list(self.__annotations.values()) 103 104 @property 105 def items(self): 106 """Returns the items in this argument pack. 107 108 Example:: 109 110 >>> vec3 = ti.types.vector(3, ti.f32) 111 >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32) 112 >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0) 113 >>> sphere.items 114 dict_items([('center', 2), ('radius', 1.0)]) 115 """ 116 return self.__entries.items() 117 118 def __getitem__(self, key): 119 ret = self.__entries[key] 120 return ret 121 122 def __setitem__(self, key, value): 123 self.__entries[key] = value 124 index = self._calc_element_true_index(list(self.__annotations).index(key)) 125 self._write_to_device(self.__annotations[key], type(value), value, index) 126 127 def _set_entries(self, value): 128 if isinstance(value, dict): 129 value = ArgPack(self.__annotations, value) 130 for k in self.keys: 131 self[k] = value[k] 132 133 @staticmethod 134 def _make_getter(key): 135 def getter(self): 136 """Get an entry from custom argument pack by name.""" 137 return self[key] 138 139 return getter 140 141 @staticmethod 142 def _make_setter(key): 143 @python_scope 144 def setter(self, value): 145 self[key] = value 146 147 return setter 148 149 def _register_members(self): 150 # https://stackoverflow.com/questions/48448074/adding-a-property-to-an-existing-object-instance 151 cls = self.__class__ 152 new_cls_name = cls.__name__ + str(cls._instance_count) 153 cls._instance_count += 1 154 properties = {k: property(cls._make_getter(k), cls._make_setter(k)) for k in self.keys} 155 self.__class__ = type(new_cls_name, (cls,), properties) 156 157 def __len__(self): 158 """Get the number of entries in a custom argument pack.""" 159 return len(self.__entries) 160 161 def __iter__(self): 162 return self.__entries.values() 163 164 def __str__(self): 165 """Python scope argument pack array print support.""" 166 if impl.inside_kernel(): 167 item_str = ", ".join([str(k) + "=" + str(v) for k, v in self.items]) 168 return f"<ti.ArgPack {item_str}>" 169 return str(self.to_dict()) 170 171 def __repr__(self): 172 return str(self.to_dict()) 173 174 def to_dict(self): 175 """Converts the ArgPack to a dictionary. 176 177 Returns: 178 Dict: The result dictionary. 179 """ 180 res_dict = { 181 k: v.to_dict() if isinstance(v, ArgPack) else v.to_list() if isinstance(v, Matrix) else v 182 for k, v in self.__entries.items() 183 } 184 return res_dict 185 186 def _calc_element_true_index(self, old_index): 187 for i in range(old_index): 188 anno = list(self.__annotations.values())[i] 189 if ( 190 isinstance(anno, sparse_matrix_builder) 191 or isinstance(anno, ndarray_type.NdarrayType) 192 or isinstance(anno, texture_type.TextureType) 193 or isinstance(anno, texture_type.RWTextureType) 194 or isinstance(anno, ndarray_type.NdarrayType) 195 ): 196 old_index -= 1 197 return old_index 198 199 def _write_to_device(self, needed, provided, v, index): 200 if isinstance(needed, ArgPackType): 201 if not isinstance(v, ArgPack): 202 raise TaichiRuntimeTypeError.get(index, str(needed), str(provided)) 203 self.__argpack.set_arg_nested_argpack(index, v.__argpack) 204 else: 205 # Note: do not use sth like "needed == f32". That would be slow. 206 if id(needed) in primitive_types.real_type_ids: 207 if not isinstance(v, (float, int, np.floating, np.integer)): 208 raise TaichiRuntimeTypeError.get(index, needed.to_string(), provided) 209 self.__argpack.set_arg_float((index,), float(v)) 210 elif id(needed) in primitive_types.integer_type_ids: 211 if not isinstance(v, (int, np.integer)): 212 raise TaichiRuntimeTypeError.get(index, needed.to_string(), provided) 213 if is_signed(cook_dtype(needed)): 214 self.__argpack.set_arg_int((index,), int(v)) 215 else: 216 self.__argpack.set_arg_uint((index,), int(v)) 217 elif isinstance(needed, sparse_matrix_builder): 218 pass 219 elif isinstance(needed, ndarray_type.NdarrayType) and isinstance(v, taichi.lang._ndarray.Ndarray): 220 pass 221 elif isinstance(needed, texture_type.TextureType) and isinstance(v, taichi.lang._texture.Texture): 222 pass 223 elif isinstance(needed, texture_type.RWTextureType) and isinstance(v, taichi.lang._texture.Texture): 224 pass 225 elif isinstance(needed, ndarray_type.NdarrayType): 226 pass 227 elif isinstance(needed, MatrixType): 228 if needed.dtype in primitive_types.real_types: 229 230 def cast_func(x): 231 if not isinstance(x, (int, float, np.integer, np.floating)): 232 raise TaichiRuntimeTypeError.get(index, needed.dtype.to_string(), type(x)) 233 return float(x) 234 235 elif needed.dtype in primitive_types.integer_types: 236 237 def cast_func(x): 238 if not isinstance(x, (int, np.integer)): 239 raise TaichiRuntimeTypeError.get(index, needed.dtype.to_string(), type(x)) 240 return int(x) 241 242 else: 243 raise ValueError(f"Matrix dtype {needed.dtype} is not integer type or real type.") 244 245 if needed.ndim == 2: 246 v = [cast_func(v[i, j]) for i in range(needed.n) for j in range(needed.m)] 247 else: 248 v = [cast_func(v[i]) for i in range(needed.n)] 249 v = needed(*v) 250 needed.set_argpack_struct_args(v, self.__argpack, (index,)) 251 elif isinstance(needed, StructType): 252 if not isinstance(v, needed): 253 raise TaichiRuntimeTypeError.get(index, str(needed), provided) 254 needed.set_argpack_struct_args(v, self.__argpack, (index,)) 255 else: 256 raise ValueError(f"Argument type mismatch. Expecting {needed}, got {type(v)}.")
The ArgPack Type Class.
The ArgPack operates as a dictionary-like data pack, storing members as (key, value) pairs. Members stored can
range from scalars and matrices to other dictionary-like structures. Distinguished from structs, ArgPack can
accommodate buffer types such as NdarrayType and TextureType from Taichi. However, unlike ti.Struct which
serves as a data container, ArgPack functions as a reference container. It's important to note that ArgPack
cannot be nested within other types except for another ArgPack, and can only be utilized as kernel parameters.
Args:
annotations (Dict[str, Union[Dict, Matrix, Struct]]): The keys and types for ArgPack members.
dtype (ArgPackType): The ArgPackType class of this ArgPack object.
entries (Dict[str, Union[Dict, Matrix, Struct]]): The keys and corresponding values for ArgPack members.
Returns:
An instance of this ArgPack.
Example::
>>> vec3 = ti.types.vector(3, ti.f32)
>>> pack_type = ti.ArgPackType(v=vec3, t=ti.f32)
>>> a = pack_type(v=vec3([0, 0, 0]), t=1.0)
>>> print(a.items)
dict_items([('v', [0. 0. 0.]), ('t', 1.0)])
57 def __init__(self, annotations, dtype, *args, **kwargs): 58 # converts dicts to argument packs 59 if len(args) == 1 and kwargs == {} and isinstance(args[0], dict): 60 self.__entries = args[0] 61 elif len(args) == 0: 62 self.__entries = kwargs 63 else: 64 raise TaichiSyntaxError( 65 "Custom argument packs need to be initialized using either dictionary or keyword arguments" 66 ) 67 if annotations.keys() != self.__entries.keys(): 68 raise TaichiSyntaxError("ArgPack annotations keys not equals to entries keys.") 69 self.__annotations = annotations 70 for k, v in self.__entries.items(): 71 self.__entries[k] = v if in_python_scope() else impl.expr_init(v) 72 self._register_members() 73 self.__dtype = dtype 74 self.__argpack = impl.get_runtime().prog.create_argpack(self.__dtype) 75 for i, (k, v) in enumerate(self.__entries.items()): 76 self._write_to_device(self.__annotations[k], type(v), v, self._calc_element_true_index(i))
82 @property 83 def keys(self): 84 """Returns the list of member names in string format. 85 86 Example:: 87 88 >>> vec3 = ti.types.vector(3, ti.f32) 89 >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32) 90 >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0) 91 >>> sphere.keys 92 ['center', 'radius'] 93 """ 94 return list(self.__entries.keys())
Returns the list of member names in string format.
Example::
>>> vec3 = ti.types.vector(3, ti.f32)
>>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32)
>>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0)
>>> sphere.keys
['center', 'radius']
104 @property 105 def items(self): 106 """Returns the items in this argument pack. 107 108 Example:: 109 110 >>> vec3 = ti.types.vector(3, ti.f32) 111 >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32) 112 >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0) 113 >>> sphere.items 114 dict_items([('center', 2), ('radius', 1.0)]) 115 """ 116 return self.__entries.items()
Returns the items in this argument pack.
Example::
>>> vec3 = ti.types.vector(3, ti.f32)
>>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32)
>>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0)
>>> sphere.items
dict_items([('center', 2), ('radius', 1.0)])
174 def to_dict(self): 175 """Converts the ArgPack to a dictionary. 176 177 Returns: 178 Dict: The result dictionary. 179 """ 180 res_dict = { 181 k: v.to_dict() if isinstance(v, ArgPack) else v.to_list() if isinstance(v, Matrix) else v 182 for k, v in self.__entries.items() 183 } 184 return res_dict
Converts the ArgPack to a dictionary.
Returns: Dict: The result dictionary.
433class BitpackedFields: 434 """Taichi bitpacked fields, where fields with quantized types are packed together. 435 436 Args: 437 max_num_bits (int): Maximum number of bits all fields inside can occupy in total. Only 32 or 64 is allowed. 438 """ 439 440 def __init__(self, max_num_bits): 441 self.fields = [] 442 self.bit_struct_type_builder = _ti_core.BitStructTypeBuilder(max_num_bits) 443 444 def place(self, *args, shared_exponent=False): 445 """Places a list of fields with quantized types inside. 446 447 Args: 448 *args (List[Field]): A list of fields with quantized types to place. 449 shared_exponent (bool): Whether the fields have a shared exponent. 450 """ 451 if shared_exponent: 452 self.bit_struct_type_builder.begin_placing_shared_exponent() 453 count = 0 454 for arg in args: 455 assert isinstance(arg, Field) 456 for var in arg._get_field_members(): 457 self.fields.append((var.ptr, self.bit_struct_type_builder.add_member(var.ptr.get_dt()))) 458 count += 1 459 if shared_exponent: 460 self.bit_struct_type_builder.end_placing_shared_exponent() 461 if count <= 1: 462 raise TaichiSyntaxError("At least 2 fields need to be placed when shared_exponent=True")
Taichi bitpacked fields, where fields with quantized types are packed together.
Args: max_num_bits (int): Maximum number of bits all fields inside can occupy in total. Only 32 or 64 is allowed.
444 def place(self, *args, shared_exponent=False): 445 """Places a list of fields with quantized types inside. 446 447 Args: 448 *args (List[Field]): A list of fields with quantized types to place. 449 shared_exponent (bool): Whether the fields have a shared exponent. 450 """ 451 if shared_exponent: 452 self.bit_struct_type_builder.begin_placing_shared_exponent() 453 count = 0 454 for arg in args: 455 assert isinstance(arg, Field) 456 for var in arg._get_field_members(): 457 self.fields.append((var.ptr, self.bit_struct_type_builder.add_member(var.ptr.get_dt()))) 458 count += 1 459 if shared_exponent: 460 self.bit_struct_type_builder.end_placing_shared_exponent() 461 if count <= 1: 462 raise TaichiSyntaxError("At least 2 fields need to be placed when shared_exponent=True")
Places a list of fields with quantized types inside.
Args: *args (List[Field]): A list of fields with quantized types to place. shared_exponent (bool): Whether the fields have a shared exponent.
21class DeviceCapability: 22 spirv_version_1_3 = "spirv_version=66304" 23 spirv_version_1_4 = "spirv_version=66560" 24 spirv_version_1_5 = "spirv_version=66816" 25 spirv_has_int8 = "spirv_has_int8" 26 spirv_has_int16 = "spirv_has_int16" 27 spirv_has_int64 = "spirv_has_int64" 28 spirv_has_float16 = "spirv_has_float16" 29 spirv_has_float64 = "spirv_has_float64" 30 spirv_has_atomic_int64 = "spirv_has_atomic_int64" 31 spirv_has_atomic_float16 = "spirv_has_atomic_float16" 32 spirv_has_atomic_float16_add = "spirv_has_atomic_float16_add" 33 spirv_has_atomic_float16_minmax = "spirv_has_atomic_float16_minmax" 34 spirv_has_atomic_float = "spirv_has_atomic_float" 35 spirv_has_atomic_float_add = "spirv_has_atomic_float_add" 36 spirv_has_atomic_float_minmax = "spirv_has_atomic_float_minmax" 37 spirv_has_atomic_float64 = "spirv_has_atomic_float64" 38 spirv_has_atomic_float64_add = "spirv_has_atomic_float64_add" 39 spirv_has_atomic_float64_minmax = "spirv_has_atomic_float64_minmax" 40 spirv_has_variable_ptr = "spirv_has_variable_ptr" 41 spirv_has_physical_storage_buffer = "spirv_has_physical_storage_buffer" 42 spirv_has_subgroup_basic = "spirv_has_subgroup_basic" 43 spirv_has_subgroup_vote = "spirv_has_subgroup_vote" 44 spirv_has_subgroup_arithmetic = "spirv_has_subgroup_arithmetic" 45 spirv_has_subgroup_ballot = "spirv_has_subgroup_ballot" 46 spirv_has_non_semantic_info = "spirv_has_non_semantic_info" 47 spirv_has_no_integer_wrap_decoration = "spirv_has_no_integer_wrap_decoration"
18class Field: 19 """Taichi field class. 20 21 A field is constructed by a list of field members. 22 For example, a scalar field has 1 field member, while a 3x3 matrix field has 9 field members. 23 A field member is a Python Expr wrapping a C++ FieldExpression. 24 25 Args: 26 vars (List[Expr]): Field members. 27 """ 28 29 def __init__(self, _vars): 30 assert all(_vars) 31 self.vars = _vars 32 self.host_accessors = None 33 self.grad = None 34 self.dual = None 35 36 @property 37 def snode(self): 38 """Gets representative SNode for info purposes. 39 40 Returns: 41 SNode: Representative SNode (SNode of first field member). 42 """ 43 return self._snode 44 45 @property 46 def _snode(self): 47 """Gets representative SNode for info purposes. 48 49 Returns: 50 SNode: Representative SNode (SNode of first field member). 51 """ 52 return taichi.lang.snode.SNode(self.vars[0].ptr.snode()) 53 54 @property 55 def shape(self): 56 """Gets field shape. 57 58 Returns: 59 Tuple[Int]: Field shape. 60 """ 61 return self._snode.shape 62 63 @property 64 def dtype(self): 65 """Gets data type of each individual value. 66 67 Returns: 68 DataType: Data type of each individual value. 69 """ 70 return self._snode._dtype 71 72 @property 73 def _name(self): 74 """Gets field name. 75 76 Returns: 77 str: Field name. 78 """ 79 return self._snode._name 80 81 def parent(self, n=1): 82 """Gets an ancestor of the representative SNode in the SNode tree. 83 84 Args: 85 n (int): the number of levels going up from the representative SNode. 86 87 Returns: 88 SNode: The n-th parent of the representative SNode. 89 """ 90 return self.snode.parent(n) 91 92 def _get_field_members(self): 93 """Gets field members. 94 95 Returns: 96 List[Expr]: Field members. 97 """ 98 return self.vars 99 100 def _loop_range(self): 101 """Gets SNode of representative field member for loop range info. 102 103 Returns: 104 taichi_python.SNode: SNode of representative (first) field member. 105 """ 106 return self.vars[0].ptr.snode() 107 108 def _set_grad(self, grad): 109 """Sets corresponding grad field (reverse mode). 110 Args: 111 grad (Field): Corresponding grad field. 112 """ 113 self.grad = grad 114 115 def _set_dual(self, dual): 116 """Sets corresponding dual field (forward mode). 117 118 Args: 119 dual (Field): Corresponding dual field. 120 """ 121 self.dual = dual 122 123 @python_scope 124 def fill(self, val): 125 """Fills `self` with a specific value. 126 127 Args: 128 val (Union[int, float]): Value to fill. 129 """ 130 raise NotImplementedError() 131 132 @python_scope 133 def to_numpy(self, dtype=None): 134 """Converts `self` to a numpy array. 135 136 Args: 137 dtype (DataType, optional): The desired data type of returned numpy array. 138 139 Returns: 140 numpy.ndarray: The result numpy array. 141 """ 142 raise NotImplementedError() 143 144 @python_scope 145 def to_torch(self, device=None): 146 """Converts `self` to a torch tensor. 147 148 Args: 149 device (torch.device, optional): The desired device of returned tensor. 150 151 Returns: 152 torch.tensor: The result torch tensor. 153 """ 154 raise NotImplementedError() 155 156 @python_scope 157 def to_paddle(self, place=None): 158 """Converts `self` to a paddle tensor. 159 160 Args: 161 place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor. 162 163 Returns: 164 paddle.Tensor: The result paddle tensor. 165 """ 166 raise NotImplementedError() 167 168 @python_scope 169 def from_numpy(self, arr): 170 """Loads all elements from a numpy array. 171 172 The shape of the numpy array needs to be the same as `self`. 173 174 Args: 175 arr (numpy.ndarray): The source numpy array. 176 """ 177 raise NotImplementedError() 178 179 @python_scope 180 def _from_external_arr(self, arr): 181 raise NotImplementedError() 182 183 @python_scope 184 def from_torch(self, arr): 185 """Loads all elements from a torch tensor. 186 187 The shape of the torch tensor needs to be the same as `self`. 188 189 Args: 190 arr (torch.tensor): The source torch tensor. 191 """ 192 self._from_external_arr(arr.contiguous()) 193 194 @python_scope 195 def from_paddle(self, arr): 196 """Loads all elements from a paddle tensor. 197 198 The shape of the paddle tensor needs to be the same as `self`. 199 200 Args: 201 arr (paddle.Tensor): The source paddle tensor. 202 """ 203 self.from_numpy(arr) 204 205 @python_scope 206 def copy_from(self, other): 207 """Copies all elements from another field. 208 209 The shape of the other field needs to be the same as `self`. 210 211 Args: 212 other (Field): The source field. 213 """ 214 if not isinstance(other, Field): 215 raise TypeError("Cannot copy from a non-field object") 216 if self.shape != other.shape: 217 raise ValueError(f"ti.field shape {self.shape} does not match" f" the source field shape {other.shape}") 218 from taichi._kernels import tensor_to_tensor # pylint: disable=C0415 219 220 tensor_to_tensor(self, other) 221 222 @python_scope 223 def __setitem__(self, key, value): 224 """Sets field element in Python scope. 225 226 Args: 227 key (Union[List[int], int, None]): Coordinates of the field element. 228 value (element type): Value to set. 229 """ 230 raise NotImplementedError() 231 232 @python_scope 233 def __getitem__(self, key): 234 """Gets field element in Python scope. 235 236 Args: 237 key (Union[List[int], int, None]): Coordinates of the field element. 238 239 Returns: 240 element type: Value retrieved. 241 """ 242 raise NotImplementedError() 243 244 def __str__(self): 245 if taichi.lang.impl.inside_kernel(): 246 return self.__repr__() # make pybind11 happy, see Matrix.__str__ 247 if self._snode.ptr is None: 248 return "<Field: Definition of this field is incomplete>" 249 return str(self.to_numpy()) 250 251 def _pad_key(self, key): 252 if key is None: 253 key = () 254 if not isinstance(key, (tuple, list)): 255 key = (key,) 256 257 if len(key) != len(self.shape): 258 raise AssertionError("Slicing is not supported on ti.field") 259 260 return key + ((0,) * (_ti_core.get_max_num_indices() - len(key))) 261 262 def _initialize_host_accessors(self): 263 if self.host_accessors: 264 return 265 taichi.lang.impl.get_runtime().materialize() 266 self.host_accessors = [SNodeHostAccessor(e.ptr.snode()) for e in self.vars] 267 268 def _host_access(self, key): 269 return [SNodeHostAccess(e, key) for e in self.host_accessors] 270 271 def __iter__(self): 272 raise NotImplementedError("Struct for is only available in Taichi scope.")
Taichi field class.
A field is constructed by a list of field members. For example, a scalar field has 1 field member, while a 3x3 matrix field has 9 field members. A field member is a Python Expr wrapping a C++ FieldExpression.
Args: vars (List[Expr]): Field members.
36 @property 37 def snode(self): 38 """Gets representative SNode for info purposes. 39 40 Returns: 41 SNode: Representative SNode (SNode of first field member). 42 """ 43 return self._snode
Gets representative SNode for info purposes.
Returns: SNode: Representative SNode (SNode of first field member).
54 @property 55 def shape(self): 56 """Gets field shape. 57 58 Returns: 59 Tuple[Int]: Field shape. 60 """ 61 return self._snode.shape
Gets field shape.
Returns: Tuple[Int]: Field shape.
63 @property 64 def dtype(self): 65 """Gets data type of each individual value. 66 67 Returns: 68 DataType: Data type of each individual value. 69 """ 70 return self._snode._dtype
Gets data type of each individual value.
Returns: DataType: Data type of each individual value.
81 def parent(self, n=1): 82 """Gets an ancestor of the representative SNode in the SNode tree. 83 84 Args: 85 n (int): the number of levels going up from the representative SNode. 86 87 Returns: 88 SNode: The n-th parent of the representative SNode. 89 """ 90 return self.snode.parent(n)
Gets an ancestor of the representative SNode in the SNode tree.
Args: n (int): the number of levels going up from the representative SNode.
Returns: SNode: The n-th parent of the representative SNode.
123 @python_scope 124 def fill(self, val): 125 """Fills `self` with a specific value. 126 127 Args: 128 val (Union[int, float]): Value to fill. 129 """ 130 raise NotImplementedError()
Fills self with a specific value.
Args: val (Union[int, float]): Value to fill.
132 @python_scope 133 def to_numpy(self, dtype=None): 134 """Converts `self` to a numpy array. 135 136 Args: 137 dtype (DataType, optional): The desired data type of returned numpy array. 138 139 Returns: 140 numpy.ndarray: The result numpy array. 141 """ 142 raise NotImplementedError()
Converts self to a numpy array.
Args: dtype (DataType, optional): The desired data type of returned numpy array.
Returns: numpy.ndarray: The result numpy array.
144 @python_scope 145 def to_torch(self, device=None): 146 """Converts `self` to a torch tensor. 147 148 Args: 149 device (torch.device, optional): The desired device of returned tensor. 150 151 Returns: 152 torch.tensor: The result torch tensor. 153 """ 154 raise NotImplementedError()
Converts self to a torch tensor.
Args: device (torch.device, optional): The desired device of returned tensor.
Returns: torch.tensor: The result torch tensor.
156 @python_scope 157 def to_paddle(self, place=None): 158 """Converts `self` to a paddle tensor. 159 160 Args: 161 place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor. 162 163 Returns: 164 paddle.Tensor: The result paddle tensor. 165 """ 166 raise NotImplementedError()
Converts self to a paddle tensor.
Args: place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor.
Returns: paddle.Tensor: The result paddle tensor.
168 @python_scope 169 def from_numpy(self, arr): 170 """Loads all elements from a numpy array. 171 172 The shape of the numpy array needs to be the same as `self`. 173 174 Args: 175 arr (numpy.ndarray): The source numpy array. 176 """ 177 raise NotImplementedError()
Loads all elements from a numpy array.
The shape of the numpy array needs to be the same as self.
Args: arr (numpy.ndarray): The source numpy array.
183 @python_scope 184 def from_torch(self, arr): 185 """Loads all elements from a torch tensor. 186 187 The shape of the torch tensor needs to be the same as `self`. 188 189 Args: 190 arr (torch.tensor): The source torch tensor. 191 """ 192 self._from_external_arr(arr.contiguous())
Loads all elements from a torch tensor.
The shape of the torch tensor needs to be the same as self.
Args: arr (torch.tensor): The source torch tensor.
194 @python_scope 195 def from_paddle(self, arr): 196 """Loads all elements from a paddle tensor. 197 198 The shape of the paddle tensor needs to be the same as `self`. 199 200 Args: 201 arr (paddle.Tensor): The source paddle tensor. 202 """ 203 self.from_numpy(arr)
Loads all elements from a paddle tensor.
The shape of the paddle tensor needs to be the same as self.
Args: arr (paddle.Tensor): The source paddle tensor.
205 @python_scope 206 def copy_from(self, other): 207 """Copies all elements from another field. 208 209 The shape of the other field needs to be the same as `self`. 210 211 Args: 212 other (Field): The source field. 213 """ 214 if not isinstance(other, Field): 215 raise TypeError("Cannot copy from a non-field object") 216 if self.shape != other.shape: 217 raise ValueError(f"ti.field shape {self.shape} does not match" f" the source field shape {other.shape}") 218 from taichi._kernels import tensor_to_tensor # pylint: disable=C0415 219 220 tensor_to_tensor(self, other)
Copies all elements from another field.
The shape of the other field needs to be the same as self.
Args: other (Field): The source field.
Members:
unknown
r8
rg8
rgba8
rgba8srgb
bgra8
bgra8srgb
r8u
rg8u
rgba8u
r8i
rg8i
rgba8i
r16
rg16
rgb16
rgba16
r16u
rg16u
rgb16u
rgba16u
r16i
rg16i
rgb16i
rgba16i
r16f
rg16f
rgb16f
rgba16f
r32u
rg32u
rgb32u
rgba32u
r32i
rg32i
rgb32i
rgba32i
r32f
rg32f
rgb32f
rgba32f
depth16
depth24stencil8
depth32f
Members:
AOS
SOA
NULL
207@_gen_swizzles 208class Matrix(TaichiOperations): 209 """The matrix class. 210 211 A matrix is a 2-D rectangular array with scalar entries, it's row-majored, and is 212 aligned continuously. We recommend only use matrix with no more than 32 elements for 213 efficiency considerations. 214 215 Note: in taichi a matrix is strictly two-dimensional and only stores scalars. 216 217 Args: 218 arr (Union[list, tuple, np.ndarray]): the initial values of a matrix. 219 dt (:mod:`~taichi.types.primitive_types`): the element data type. 220 ndim (int optional): the number of dimensions of the matrix; forced reshape if given. 221 222 Example:: 223 224 use a 2d list to initialize a matrix 225 226 >>> @ti.kernel 227 >>> def test(): 228 >>> n = 5 229 >>> M = ti.Matrix([[0] * n for _ in range(n)], ti.i32) 230 >>> print(M) # a 5x5 matrix with integer elements 231 232 get the number of rows and columns via the `n`, `m` property: 233 234 >>> M = ti.Matrix([[0, 1], [2, 3], [4, 5]], ti.i32) 235 >>> M.n # number of rows 236 3 237 >>> M.m # number of cols 238 >>> 2 239 240 you can even initialize a matrix with an empty list: 241 242 >>> M = ti.Matrix([[], []], ti.i32) 243 >>> M.n 244 2 245 >>> M.m 246 0 247 """ 248 249 _is_taichi_class = True 250 _is_matrix_class = True 251 __array_priority__ = 1000 252 253 def __init__(self, arr, dt=None): 254 if not isinstance(arr, (list, tuple, np.ndarray)): 255 raise TaichiTypeError("An Matrix/Vector can only be initialized with an array-like object") 256 if len(arr) == 0: 257 self.ndim = 0 258 self.n, self.m = 0, 0 259 self.entries = np.array([]) 260 self.is_host_access = False 261 elif isinstance(arr[0], Matrix): 262 raise Exception("cols/rows required when using list of vectors") 263 elif isinstance(arr[0], Iterable): # matrix 264 self.ndim = 2 265 self.n, self.m = len(arr), len(arr[0]) 266 if isinstance(arr[0][0], (SNodeHostAccess, NdarrayHostAccess)): 267 self.entries = arr 268 self.is_host_access = True 269 else: 270 self.entries = np.array(arr, None if dt is None else to_numpy_type(dt)) 271 self.is_host_access = False 272 else: # vector 273 self.ndim = 1 274 self.n, self.m = len(arr), 1 275 if isinstance(arr[0], (SNodeHostAccess, NdarrayHostAccess)): 276 self.entries = arr 277 self.is_host_access = True 278 else: 279 self.entries = np.array(arr, None if dt is None else to_numpy_type(dt)) 280 self.is_host_access = False 281 282 if self.n * self.m > 32: 283 warning( 284 f"Taichi matrices/vectors with {self.n}x{self.m} > 32 entries are not suggested." 285 " Matrices/vectors will be automatically unrolled at compile-time for performance." 286 " So the compilation time could be extremely long if the matrix size is too big." 287 " You may use a field to store a large matrix like this, e.g.:\n" 288 f" x = ti.field(ti.f32, ({self.n}, {self.m})).\n" 289 " See https://docs.taichi-lang.org/docs/field#matrix-size" 290 " for more details.", 291 UserWarning, 292 stacklevel=2, 293 ) 294 295 def get_shape(self): 296 if self.ndim == 1: 297 return (self.n,) 298 if self.ndim == 2: 299 return (self.n, self.m) 300 return None 301 302 def __matmul__(self, other): 303 """Matrix-matrix or matrix-vector multiply. 304 305 Args: 306 other (Union[Matrix, Vector]): a matrix or a vector. 307 308 Returns: 309 The matrix-matrix product or matrix-vector product. 310 311 """ 312 from taichi.lang import matrix_ops # pylint: disable=C0415 313 314 return matrix_ops.matmul(self, other) 315 316 # host access & python scope operation 317 def __len__(self): 318 """Get the length of each row of a matrix""" 319 # TODO: When this is a vector, should return its dimension? 320 return self.n 321 322 def __iter__(self): 323 if self.ndim == 1: 324 return (self[i] for i in range(self.n)) 325 return ([self[i, j] for j in range(self.m)] for i in range(self.n)) 326 327 def __getitem__(self, indices): 328 """Access to the element at the given indices in a matrix. 329 330 Args: 331 indices (Sequence[Expr]): the indices of the element. 332 333 Returns: 334 The value of the element at a specific position of a matrix. 335 336 """ 337 entry = self._get_entry(indices) 338 if self.is_host_access: 339 return _read_host_access(entry) 340 return entry 341 342 @python_scope 343 def __setitem__(self, indices, item): 344 """Set the element value at the given indices in a matrix. 345 346 Args: 347 indices (Sequence[Expr]): the indices of a element. 348 349 """ 350 if self.is_host_access: 351 entry = self._get_entry(indices) 352 _write_host_access(entry, item) 353 else: 354 if not isinstance(indices, (list, tuple)): 355 indices = [indices] 356 assert len(indices) in [1, 2] 357 assert len(indices) == self.ndim, f"Expected {self.ndim} indices, got {len(indices)}" 358 if self.ndim == 1: 359 self.entries[indices[0]] = item 360 else: 361 self.entries[indices[0]][indices[1]] = item 362 363 def _get_entry(self, indices): 364 if not isinstance(indices, (list, tuple)): 365 indices = [indices] 366 assert len(indices) in [1, 2] 367 assert len(indices) == self.ndim, f"Expected {self.ndim} indices, got {len(indices)}" 368 if self.ndim == 1: 369 return self.entries[indices[0]] 370 return self.entries[indices[0]][indices[1]] 371 372 def _get_slice(self, a, b): 373 if isinstance(a, slice): 374 a = range(a.start or 0, a.stop or self.n, a.step or 1) 375 if isinstance(b, slice): 376 b = range(b.start or 0, b.stop or self.m, b.step or 1) 377 if isinstance(a, range) and isinstance(b, range): 378 return Matrix([[self._get_entry(i, j) for j in b] for i in a]) 379 if isinstance(a, range): # b is not range 380 return Vector([self._get_entry(i, b) for i in a]) 381 # a is not range while b is range 382 return Vector([self._get_entry(a, j) for j in b]) 383 384 @python_scope 385 def _set_entries(self, value): 386 if isinstance(value, Matrix): 387 value = value.to_list() 388 if self.is_host_access: 389 if self.ndim == 1: 390 for i in range(self.n): 391 _write_host_access(self.entries[i], value[i]) 392 else: 393 for i in range(self.n): 394 for j in range(self.m): 395 _write_host_access(self.entries[i][j], value[i][j]) 396 else: 397 if self.ndim == 1: 398 for i in range(self.n): 399 self.entries[i] = value[i] 400 else: 401 for i in range(self.n): 402 for j in range(self.m): 403 self.entries[i][j] = value[i][j] 404 405 @property 406 def _members(self): 407 return self.entries 408 409 def to_list(self): 410 """Return this matrix as a 1D `list`. 411 412 This is similar to `numpy.ndarray`'s `flatten` and `ravel` methods, 413 the difference is that this function always returns a new list. 414 """ 415 if self.is_host_access: 416 if self.ndim == 1: 417 return [_read_host_access(self.entries[i]) for i in range(self.n)] 418 assert self.ndim == 2 419 return [[_read_host_access(self.entries[i][j]) for j in range(self.m)] for i in range(self.n)] 420 return self.entries.tolist() 421 422 @taichi_scope 423 def cast(self, dtype): 424 """Cast the matrix elements to a specified data type. 425 426 Args: 427 dtype (:mod:`~taichi.types.primitive_types`): data type of the 428 returned matrix. 429 430 Returns: 431 :class:`taichi.Matrix`: A new matrix with the specified data dtype. 432 433 Example:: 434 435 >>> A = ti.Matrix([0, 1, 2], ti.i32) 436 >>> B = A.cast(ti.f32) 437 >>> B 438 [0.0, 1.0, 2.0] 439 """ 440 if self.ndim == 1: 441 return Vector([ops_mod.cast(self[i], dtype) for i in range(self.n)]) 442 return Matrix([[ops_mod.cast(self[i, j], dtype) for j in range(self.m)] for i in range(self.n)]) 443 444 def trace(self): 445 """The sum of a matrix diagonal elements. 446 447 To call this method the matrix must be square-like. 448 449 Returns: 450 The sum of a matrix diagonal elements. 451 452 Example:: 453 454 >>> m = ti.Matrix([[1, 2], [3, 4]]) 455 >>> m.trace() 456 5 457 """ 458 # pylint: disable-msg=C0415 459 from taichi.lang import matrix_ops 460 461 return matrix_ops.trace(self) 462 463 def inverse(self): 464 """Returns the inverse of this matrix. 465 466 Note: 467 The matrix dimension should be less than or equal to 4. 468 469 Returns: 470 :class:`~taichi.Matrix`: The inverse of a matrix. 471 472 Raises: 473 Exception: Inversions of matrices with sizes >= 5 are not supported. 474 """ 475 from taichi.lang import matrix_ops # pylint: disable=C0415 476 477 return matrix_ops.inverse(self) 478 479 def normalized(self, eps=0): 480 """Normalize a vector, i.e. matrices with the second dimension being 481 equal to one. 482 483 The normalization of a vector `v` is a vector of length 1 484 and has the same direction with `v`. It's equal to `v/|v|`. 485 486 Args: 487 eps (float): a safe-guard value for sqrt, usually 0. 488 489 Example:: 490 491 >>> a = ti.Vector([3, 4], ti.f32) 492 >>> a.normalized() 493 [0.6, 0.8] 494 """ 495 # pylint: disable-msg=C0415 496 from taichi.lang import matrix_ops 497 498 return matrix_ops.normalized(self, eps) 499 500 def transpose(self): 501 """Returns the transpose of a matrix. 502 503 Returns: 504 :class:`~taichi.Matrix`: The transpose of this matrix. 505 506 Example:: 507 508 >>> A = ti.Matrix([[0, 1], [2, 3]]) 509 >>> A.transpose() 510 [[0, 2], [1, 3]] 511 """ 512 # pylint: disable=C0415 513 from taichi.lang import matrix_ops 514 515 return matrix_ops.transpose(self) 516 517 @taichi_scope 518 def determinant(a): 519 """Returns the determinant of this matrix. 520 521 Note: 522 The matrix dimension should be less than or equal to 4. 523 524 Returns: 525 dtype: The determinant of this matrix. 526 527 Raises: 528 Exception: Determinants of matrices with sizes >= 5 are not supported. 529 """ 530 # pylint: disable=C0415 531 from taichi.lang import matrix_ops 532 533 return matrix_ops.determinant(a) 534 535 @staticmethod 536 def diag(dim, val): 537 """Returns a diagonal square matrix with the diagonals filled 538 with `val`. 539 540 Args: 541 dim (int): the dimension of the wanted square matrix. 542 val (TypeVar): value for the diagonal elements. 543 544 Returns: 545 :class:`~taichi.Matrix`: The wanted diagonal matrix. 546 547 Example:: 548 549 >>> m = ti.Matrix.diag(3, 1) 550 [[1, 0, 0], 551 [0, 1, 0], 552 [0, 0, 1]] 553 """ 554 # pylint: disable=C0415 555 from taichi.lang import matrix_ops 556 557 return matrix_ops.diag(dim, val) 558 559 def sum(self): 560 """Return the sum of all elements. 561 562 Example:: 563 564 >>> m = ti.Matrix([[1, 2], [3, 4]]) 565 >>> m.sum() 566 10 567 """ 568 # pylint: disable=C0415 569 from taichi.lang import matrix_ops 570 571 return matrix_ops.sum(self) 572 573 def norm(self, eps=0): 574 """Returns the square root of the sum of the absolute squares 575 of its elements. 576 577 Args: 578 eps (Number): a safe-guard value for sqrt, usually 0. 579 580 Example:: 581 582 >>> a = ti.Vector([3, 4]) 583 >>> a.norm() 584 5 585 586 Returns: 587 The square root of the sum of the absolute squares of its elements. 588 """ 589 # pylint: disable=C0415 590 from taichi.lang import matrix_ops 591 592 return matrix_ops.norm(self, eps=eps) 593 594 def norm_inv(self, eps=0): 595 """The inverse of the matrix :func:`~taichi.lang.matrix.Matrix.norm`. 596 597 Args: 598 eps (float): a safe-guard value for sqrt, usually 0. 599 600 Returns: 601 The inverse of the matrix/vector `norm`. 602 """ 603 # pylint: disable=C0415 604 from taichi.lang import matrix_ops 605 606 return matrix_ops.norm_inv(self, eps=eps) 607 608 def norm_sqr(self): 609 """Returns the sum of the absolute squares of its elements.""" 610 # pylint: disable=C0415 611 from taichi.lang import matrix_ops 612 613 return matrix_ops.norm_sqr(self) 614 615 def max(self): 616 """Returns the maximum element value.""" 617 # pylint: disable=C0415 618 from taichi.lang import matrix_ops 619 620 return matrix_ops.max(self) 621 622 def min(self): 623 """Returns the minimum element value.""" 624 # pylint: disable=C0415 625 from taichi.lang import matrix_ops 626 627 return matrix_ops.min(self) 628 629 def any(self): 630 """Test whether any element not equal zero. 631 632 Returns: 633 bool: `True` if any element is not equal zero, `False` otherwise. 634 635 Example:: 636 637 >>> v = ti.Vector([0, 0, 1]) 638 >>> v.any() 639 True 640 """ 641 # pylint: disable=C0415 642 from taichi.lang import matrix_ops 643 644 return matrix_ops.any(self) 645 646 def all(self): 647 """Test whether all element not equal zero. 648 649 Returns: 650 bool: `True` if all elements are not equal zero, `False` otherwise. 651 652 Example:: 653 654 >>> v = ti.Vector([0, 0, 1]) 655 >>> v.all() 656 False 657 """ 658 # pylint: disable=C0415 659 from taichi.lang import matrix_ops 660 661 return matrix_ops.all(self) 662 663 def fill(self, val): 664 """Fills the matrix with a specified value. 665 666 Args: 667 val (Union[int, float]): Value to fill. 668 669 Example:: 670 671 >>> A = ti.Matrix([0, 1, 2, 3]) 672 >>> A.fill(-1) 673 >>> A 674 [-1, -1, -1, -1] 675 """ 676 # pylint: disable=C0415 677 from taichi.lang import matrix_ops 678 679 return matrix_ops.fill(self, val) 680 681 def to_numpy(self): 682 """Converts this matrix to a numpy array. 683 684 Returns: 685 numpy.ndarray: The result numpy array. 686 687 Example:: 688 689 >>> A = ti.Matrix([[0], [1], [2], [3]]) 690 >>> A.to_numpy() 691 >>> A 692 array([[0], [1], [2], [3]]) 693 """ 694 if self.is_host_access: 695 return np.array(self.to_list()) 696 return self.entries 697 698 @taichi_scope 699 def __ti_repr__(self): 700 yield "[" 701 for i in range(self.n): 702 if i: 703 yield ", " 704 if self.m != 1: 705 yield "[" 706 for j in range(self.m): 707 if j: 708 yield ", " 709 yield self(i, j) 710 if self.m != 1: 711 yield "]" 712 yield "]" 713 714 def __str__(self): 715 """Python scope matrix print support.""" 716 if impl.inside_kernel(): 717 """ 718 It seems that when pybind11 got an type mismatch, it will try 719 to invoke `repr` to show the object... e.g.: 720 721 TypeError: make_const_expr_f32(): incompatible function arguments. The following argument types are supported: 722 1. (arg0: float) -> taichi_python.Expr 723 724 Invoked with: <Taichi 2x1 Matrix> 725 726 So we have to make it happy with a dummy string... 727 """ 728 return f"<{self.n}x{self.m} ti.Matrix>" 729 return str(self.to_numpy()) 730 731 def __repr__(self): 732 return str(self.to_numpy()) 733 734 @staticmethod 735 @taichi_scope 736 def zero(dt, n, m=None): 737 """Constructs a Matrix filled with zeros. 738 739 Args: 740 dt (DataType): The desired data type. 741 n (int): The first dimension (row) of the matrix. 742 m (int, optional): The second dimension (column) of the matrix. 743 744 Returns: 745 :class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance filled with zeros. 746 747 """ 748 from taichi.lang import matrix_ops # pylint: disable=C0415 749 750 if m is None: 751 return matrix_ops._filled_vector(n, dt, 0) 752 return matrix_ops._filled_matrix(n, m, dt, 0) 753 754 @staticmethod 755 @taichi_scope 756 def one(dt, n, m=None): 757 """Constructs a Matrix filled with ones. 758 759 Args: 760 dt (DataType): The desired data type. 761 n (int): The first dimension (row) of the matrix. 762 m (int, optional): The second dimension (column) of the matrix. 763 764 Returns: 765 :class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance filled with ones. 766 767 """ 768 from taichi.lang import matrix_ops # pylint: disable=C0415 769 770 if m is None: 771 return matrix_ops._filled_vector(n, dt, 1) 772 return matrix_ops._filled_matrix(n, m, dt, 1) 773 774 @staticmethod 775 @taichi_scope 776 def unit(n, i, dt=None): 777 """Constructs a n-D vector with the `i`-th entry being equal to one and 778 the remaining entries are all zeros. 779 780 Args: 781 n (int): The length of the vector. 782 i (int): The index of the entry that will be filled with one. 783 dt (:mod:`~taichi.types.primitive_types`, optional): The desired data type. 784 785 Returns: 786 :class:`~taichi.Matrix`: The returned vector. 787 788 Example:: 789 790 >>> A = ti.Matrix.unit(3, 1) 791 >>> A 792 [0, 1, 0] 793 """ 794 from taichi.lang import matrix_ops # pylint: disable=C0415 795 796 if dt is None: 797 dt = int 798 assert 0 <= i < n 799 return matrix_ops._unit_vector(n, i, dt) 800 801 @staticmethod 802 @taichi_scope 803 def identity(dt, n): 804 """Constructs an identity Matrix with shape (n, n). 805 806 Args: 807 dt (DataType): The desired data type. 808 n (int): The number of rows/columns. 809 810 Returns: 811 :class:`~taichi.Matrix`: An `n x n` identity matrix. 812 """ 813 from taichi.lang import matrix_ops # pylint: disable=C0415 814 815 return matrix_ops._identity_matrix(n, dt) 816 817 @classmethod 818 @python_scope 819 def field( 820 cls, 821 n, 822 m, 823 dtype, 824 shape=None, 825 order=None, 826 name="", 827 offset=None, 828 needs_grad=False, 829 needs_dual=False, 830 layout=Layout.AOS, 831 ndim=None, 832 ): 833 """Construct a data container to hold all elements of the Matrix. 834 835 Args: 836 n (int): The desired number of rows of the Matrix. 837 m (int): The desired number of columns of the Matrix. 838 dtype (DataType, optional): The desired data type of the Matrix. 839 shape (Union[int, tuple of int], optional): The desired shape of the Matrix. 840 order (str, optional): order of the shape laid out in memory. 841 name (string, optional): The custom name of the field. 842 offset (Union[int, tuple of int], optional): The coordinate offset 843 of all elements in a field. 844 needs_grad (bool, optional): Whether the Matrix need grad field (reverse mode autodiff). 845 needs_dual (bool, optional): Whether the Matrix need dual field (forward mode autodiff). 846 layout (Layout, optional): The field layout, either Array Of 847 Structure (AOS) or Structure Of Array (SOA). 848 849 Returns: 850 :class:`~taichi.Matrix`: A matrix. 851 """ 852 entries = [] 853 element_dim = ndim if ndim is not None else 2 854 if isinstance(dtype, (list, tuple, np.ndarray)): 855 # set different dtype for each element in Matrix 856 # see #2135 857 if m == 1: 858 assert ( 859 len(np.shape(dtype)) == 1 and len(dtype) == n 860 ), f"Please set correct dtype list for Vector. The shape of dtype list should be ({n}, ) instead of {np.shape(dtype)}" 861 for i in range(n): 862 entries.append( 863 impl.create_field_member( 864 dtype[i], 865 name=name, 866 needs_grad=needs_grad, 867 needs_dual=needs_dual, 868 ) 869 ) 870 else: 871 assert ( 872 len(np.shape(dtype)) == 2 and len(dtype) == n and len(dtype[0]) == m 873 ), f"Please set correct dtype list for Matrix. The shape of dtype list should be ({n}, {m}) instead of {np.shape(dtype)}" 874 for i in range(n): 875 for j in range(m): 876 entries.append( 877 impl.create_field_member( 878 dtype[i][j], 879 name=name, 880 needs_grad=needs_grad, 881 needs_dual=needs_dual, 882 ) 883 ) 884 else: 885 for _ in range(n * m): 886 entries.append(impl.create_field_member(dtype, name=name, needs_grad=needs_grad, needs_dual=needs_dual)) 887 entries, entries_grad, entries_dual = zip(*entries) 888 889 entries = MatrixField(entries, n, m, element_dim) 890 if all(entries_grad): 891 entries_grad = MatrixField(entries_grad, n, m, element_dim) 892 entries._set_grad(entries_grad) 893 if all(entries_dual): 894 entries_dual = MatrixField(entries_dual, n, m, element_dim) 895 entries._set_dual(entries_dual) 896 897 impl.get_runtime().matrix_fields.append(entries) 898 899 if shape is None: 900 if offset is not None: 901 raise TaichiSyntaxError("shape cannot be None when offset is set") 902 if order is not None: 903 raise TaichiSyntaxError("shape cannot be None when order is set") 904 else: 905 if isinstance(shape, numbers.Number): 906 shape = (shape,) 907 if isinstance(offset, numbers.Number): 908 offset = (offset,) 909 dim = len(shape) 910 if offset is not None and dim != len(offset): 911 raise TaichiSyntaxError( 912 f"The dimensionality of shape and offset must be the same ({dim} != {len(offset)})" 913 ) 914 axis_seq = [] 915 shape_seq = [] 916 if order is not None: 917 if dim != len(order): 918 raise TaichiSyntaxError( 919 f"The dimensionality of shape and order must be the same ({dim} != {len(order)})" 920 ) 921 if dim != len(set(order)): 922 raise TaichiSyntaxError("The axes in order must be different") 923 for ch in order: 924 axis = ord(ch) - ord("i") 925 if axis < 0 or axis >= dim: 926 raise TaichiSyntaxError(f"Invalid axis {ch}") 927 axis_seq.append(axis) 928 shape_seq.append(shape[axis]) 929 else: 930 axis_seq = list(range(dim)) 931 shape_seq = list(shape) 932 same_level = order is None 933 if layout == Layout.SOA: 934 for e in entries._get_field_members(): 935 impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset) 936 if needs_grad: 937 for e in entries_grad._get_field_members(): 938 impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset) 939 if needs_dual: 940 for e in entries_dual._get_field_members(): 941 impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset) 942 else: 943 impl._create_snode(axis_seq, shape_seq, same_level).place(entries, offset=offset) 944 if needs_grad: 945 impl._create_snode(axis_seq, shape_seq, same_level).place(entries_grad, offset=offset) 946 if needs_dual: 947 impl._create_snode(axis_seq, shape_seq, same_level).place(entries_dual, offset=offset) 948 return entries 949 950 @classmethod 951 @python_scope 952 def ndarray(cls, n, m, dtype, shape): 953 """Defines a Taichi ndarray with matrix elements. 954 This function must be called in Python scope, and after `ti.init` is called. 955 956 Args: 957 n (int): Number of rows of the matrix. 958 m (int): Number of columns of the matrix. 959 dtype (DataType): Data type of each value. 960 shape (Union[int, tuple[int]]): Shape of the ndarray. 961 962 Example:: 963 964 The code below shows how a Taichi ndarray with matrix elements \ 965 can be declared and defined:: 966 967 >>> x = ti.Matrix.ndarray(4, 5, ti.f32, shape=(16, 8)) 968 """ 969 if isinstance(shape, numbers.Number): 970 shape = (shape,) 971 return MatrixNdarray(n, m, dtype, shape) 972 973 @staticmethod 974 def rows(rows): 975 """Constructs a matrix by concatenating a list of 976 vectors/lists row by row. Must be called in Taichi scope. 977 978 Args: 979 rows (List): A list of Vector (1-D Matrix) or a list of list. 980 981 Returns: 982 :class:`~taichi.Matrix`: A matrix. 983 984 Example:: 985 986 >>> @ti.kernel 987 >>> def test(): 988 >>> v1 = ti.Vector([1, 2, 3]) 989 >>> v2 = ti.Vector([4, 5, 6]) 990 >>> m = ti.Matrix.rows([v1, v2]) 991 >>> print(m) 992 >>> 993 >>> test() 994 [[1, 2, 3], [4, 5, 6]] 995 """ 996 from taichi.lang import matrix_ops # pylint: disable=C0415 997 998 return matrix_ops.rows(rows) 999 1000 @staticmethod 1001 def cols(cols): 1002 """Constructs a Matrix instance by concatenating Vectors/lists column by column. 1003 1004 Args: 1005 cols (List): A list of Vector (1-D Matrix) or a list of list. 1006 1007 Returns: 1008 :class:`~taichi.Matrix`: A matrix. 1009 1010 Example:: 1011 1012 >>> @ti.kernel 1013 >>> def test(): 1014 >>> v1 = ti.Vector([1, 2, 3]) 1015 >>> v2 = ti.Vector([4, 5, 6]) 1016 >>> m = ti.Matrix.cols([v1, v2]) 1017 >>> print(m) 1018 >>> 1019 >>> test() 1020 [[1, 4], [2, 5], [3, 6]] 1021 """ 1022 from taichi.lang import matrix_ops # pylint: disable=C0415 1023 1024 return matrix_ops.cols(cols) 1025 1026 def __hash__(self): 1027 # TODO: refactor KernelTemplateMapper 1028 # If not, we get `unhashable type: Matrix` when 1029 # using matrices as template arguments. 1030 return id(self) 1031 1032 def dot(self, other): 1033 """Performs the dot product of two vectors. 1034 1035 To call this method, both multiplicatives must be vectors. 1036 1037 Args: 1038 other (:class:`~taichi.Matrix`): The input Vector. 1039 1040 Returns: 1041 DataType: The dot product result (scalar) of the two Vectors. 1042 1043 Example:: 1044 1045 >>> v1 = ti.Vector([1, 2, 3]) 1046 >>> v2 = ti.Vector([3, 4, 5]) 1047 >>> v1.dot(v2) 1048 26 1049 """ 1050 from taichi.lang import matrix_ops # pylint: disable=C0415 1051 1052 return matrix_ops.dot(self, other) 1053 1054 def cross(self, other): 1055 """Performs the cross product with the input vector (1-D Matrix). 1056 1057 Both two vectors must have the same dimension <= 3. 1058 1059 For two 2d vectors (x1, y1) and (x2, y2), the return value is the 1060 scalar `x1*y2 - x2*y1`. 1061 1062 For two 3d vectors `v` and `w`, the return value is the 3d vector 1063 `v x w`. 1064 1065 Args: 1066 other (:class:`~taichi.Matrix`): The input Vector. 1067 1068 Returns: 1069 :class:`~taichi.Matrix`: The cross product of the two Vectors. 1070 """ 1071 from taichi.lang import matrix_ops # pylint: disable=C0415 1072 1073 return matrix_ops.cross(self, other) 1074 1075 def outer_product(self, other): 1076 """Performs the outer product with the input Vector (1-D Matrix). 1077 1078 The outer_product of two vectors `v = (x1, x2, ..., xn)`, 1079 `w = (y1, y2, ..., yn)` is a `n` times `n` square matrix, and its `(i, j)` 1080 entry is equal to `xi*yj`. 1081 1082 Args: 1083 other (:class:`~taichi.Matrix`): The input Vector. 1084 1085 Returns: 1086 :class:`~taichi.Matrix`: The outer product of the two Vectors. 1087 """ 1088 from taichi.lang import matrix_ops # pylint: disable=C0415 1089 1090 return matrix_ops.outer_product(self, other)
The matrix class.
A matrix is a 2-D rectangular array with scalar entries, it's row-majored, and is aligned continuously. We recommend only use matrix with no more than 32 elements for efficiency considerations.
Note: in taichi a matrix is strictly two-dimensional and only stores scalars.
Args:
arr (Union[list, tuple, np.ndarray]): the initial values of a matrix.
dt (~taichi.types.primitive_types): the element data type.
ndim (int optional): the number of dimensions of the matrix; forced reshape if given.
Example::
use a 2d list to initialize a matrix
>>> @ti.kernel
>>> def test():
>>> n = 5
>>> M = ti.Matrix([[0] * n for _ in range(n)], ti.i32)
>>> print(M) # a 5x5 matrix with integer elements
get the number of rows and columns via the `n`, `m` property:
>>> M = ti.Matrix([[0, 1], [2, 3], [4, 5]], ti.i32)
>>> M.n # number of rows
3
>>> M.m # number of cols
>>> 2
you can even initialize a matrix with an empty list:
>>> M = ti.Matrix([[], []], ti.i32)
>>> M.n
2
>>> M.m
0
253 def __init__(self, arr, dt=None): 254 if not isinstance(arr, (list, tuple, np.ndarray)): 255 raise TaichiTypeError("An Matrix/Vector can only be initialized with an array-like object") 256 if len(arr) == 0: 257 self.ndim = 0 258 self.n, self.m = 0, 0 259 self.entries = np.array([]) 260 self.is_host_access = False 261 elif isinstance(arr[0], Matrix): 262 raise Exception("cols/rows required when using list of vectors") 263 elif isinstance(arr[0], Iterable): # matrix 264 self.ndim = 2 265 self.n, self.m = len(arr), len(arr[0]) 266 if isinstance(arr[0][0], (SNodeHostAccess, NdarrayHostAccess)): 267 self.entries = arr 268 self.is_host_access = True 269 else: 270 self.entries = np.array(arr, None if dt is None else to_numpy_type(dt)) 271 self.is_host_access = False 272 else: # vector 273 self.ndim = 1 274 self.n, self.m = len(arr), 1 275 if isinstance(arr[0], (SNodeHostAccess, NdarrayHostAccess)): 276 self.entries = arr 277 self.is_host_access = True 278 else: 279 self.entries = np.array(arr, None if dt is None else to_numpy_type(dt)) 280 self.is_host_access = False 281 282 if self.n * self.m > 32: 283 warning( 284 f"Taichi matrices/vectors with {self.n}x{self.m} > 32 entries are not suggested." 285 " Matrices/vectors will be automatically unrolled at compile-time for performance." 286 " So the compilation time could be extremely long if the matrix size is too big." 287 " You may use a field to store a large matrix like this, e.g.:\n" 288 f" x = ti.field(ti.f32, ({self.n}, {self.m})).\n" 289 " See https://docs.taichi-lang.org/docs/field#matrix-size" 290 " for more details.", 291 UserWarning, 292 stacklevel=2, 293 )
409 def to_list(self): 410 """Return this matrix as a 1D `list`. 411 412 This is similar to `numpy.ndarray`'s `flatten` and `ravel` methods, 413 the difference is that this function always returns a new list. 414 """ 415 if self.is_host_access: 416 if self.ndim == 1: 417 return [_read_host_access(self.entries[i]) for i in range(self.n)] 418 assert self.ndim == 2 419 return [[_read_host_access(self.entries[i][j]) for j in range(self.m)] for i in range(self.n)] 420 return self.entries.tolist()
Return this matrix as a 1D list.
This is similar to numpy.ndarray's flatten and ravel methods,
the difference is that this function always returns a new list.
422 @taichi_scope 423 def cast(self, dtype): 424 """Cast the matrix elements to a specified data type. 425 426 Args: 427 dtype (:mod:`~taichi.types.primitive_types`): data type of the 428 returned matrix. 429 430 Returns: 431 :class:`taichi.Matrix`: A new matrix with the specified data dtype. 432 433 Example:: 434 435 >>> A = ti.Matrix([0, 1, 2], ti.i32) 436 >>> B = A.cast(ti.f32) 437 >>> B 438 [0.0, 1.0, 2.0] 439 """ 440 if self.ndim == 1: 441 return Vector([ops_mod.cast(self[i], dtype) for i in range(self.n)]) 442 return Matrix([[ops_mod.cast(self[i, j], dtype) for j in range(self.m)] for i in range(self.n)])
Cast the matrix elements to a specified data type.
Args:
dtype (~taichi.types.primitive_types): data type of the
returned matrix.
Returns:
taichi.Matrix: A new matrix with the specified data dtype.
Example::
>>> A = ti.Matrix([0, 1, 2], ti.i32)
>>> B = A.cast(ti.f32)
>>> B
[0.0, 1.0, 2.0]
444 def trace(self): 445 """The sum of a matrix diagonal elements. 446 447 To call this method the matrix must be square-like. 448 449 Returns: 450 The sum of a matrix diagonal elements. 451 452 Example:: 453 454 >>> m = ti.Matrix([[1, 2], [3, 4]]) 455 >>> m.trace() 456 5 457 """ 458 # pylint: disable-msg=C0415 459 from taichi.lang import matrix_ops 460 461 return matrix_ops.trace(self)
The sum of a matrix diagonal elements.
To call this method the matrix must be square-like.
Returns: The sum of a matrix diagonal elements.
Example::
>>> m = ti.Matrix([[1, 2], [3, 4]])
>>> m.trace()
5
463 def inverse(self): 464 """Returns the inverse of this matrix. 465 466 Note: 467 The matrix dimension should be less than or equal to 4. 468 469 Returns: 470 :class:`~taichi.Matrix`: The inverse of a matrix. 471 472 Raises: 473 Exception: Inversions of matrices with sizes >= 5 are not supported. 474 """ 475 from taichi.lang import matrix_ops # pylint: disable=C0415 476 477 return matrix_ops.inverse(self)
Returns the inverse of this matrix.
Note: The matrix dimension should be less than or equal to 4.
Returns:
~taichi.Matrix: The inverse of a matrix.
Raises: Exception: Inversions of matrices with sizes >= 5 are not supported.
479 def normalized(self, eps=0): 480 """Normalize a vector, i.e. matrices with the second dimension being 481 equal to one. 482 483 The normalization of a vector `v` is a vector of length 1 484 and has the same direction with `v`. It's equal to `v/|v|`. 485 486 Args: 487 eps (float): a safe-guard value for sqrt, usually 0. 488 489 Example:: 490 491 >>> a = ti.Vector([3, 4], ti.f32) 492 >>> a.normalized() 493 [0.6, 0.8] 494 """ 495 # pylint: disable-msg=C0415 496 from taichi.lang import matrix_ops 497 498 return matrix_ops.normalized(self, eps)
Normalize a vector, i.e. matrices with the second dimension being equal to one.
The normalization of a vector v is a vector of length 1
and has the same direction with v. It's equal to v/|v|.
Args: eps (float): a safe-guard value for sqrt, usually 0.
Example::
>>> a = ti.Vector([3, 4], ti.f32)
>>> a.normalized()
[0.6, 0.8]
500 def transpose(self): 501 """Returns the transpose of a matrix. 502 503 Returns: 504 :class:`~taichi.Matrix`: The transpose of this matrix. 505 506 Example:: 507 508 >>> A = ti.Matrix([[0, 1], [2, 3]]) 509 >>> A.transpose() 510 [[0, 2], [1, 3]] 511 """ 512 # pylint: disable=C0415 513 from taichi.lang import matrix_ops 514 515 return matrix_ops.transpose(self)
Returns the transpose of a matrix.
Returns:
~taichi.Matrix: The transpose of this matrix.
Example::
>>> A = ti.Matrix([[0, 1], [2, 3]])
>>> A.transpose()
[[0, 2], [1, 3]]
517 @taichi_scope 518 def determinant(a): 519 """Returns the determinant of this matrix. 520 521 Note: 522 The matrix dimension should be less than or equal to 4. 523 524 Returns: 525 dtype: The determinant of this matrix. 526 527 Raises: 528 Exception: Determinants of matrices with sizes >= 5 are not supported. 529 """ 530 # pylint: disable=C0415 531 from taichi.lang import matrix_ops 532 533 return matrix_ops.determinant(a)
Returns the determinant of this matrix.
Note: The matrix dimension should be less than or equal to 4.
Returns: dtype: The determinant of this matrix.
Raises: Exception: Determinants of matrices with sizes >= 5 are not supported.
535 @staticmethod 536 def diag(dim, val): 537 """Returns a diagonal square matrix with the diagonals filled 538 with `val`. 539 540 Args: 541 dim (int): the dimension of the wanted square matrix. 542 val (TypeVar): value for the diagonal elements. 543 544 Returns: 545 :class:`~taichi.Matrix`: The wanted diagonal matrix. 546 547 Example:: 548 549 >>> m = ti.Matrix.diag(3, 1) 550 [[1, 0, 0], 551 [0, 1, 0], 552 [0, 0, 1]] 553 """ 554 # pylint: disable=C0415 555 from taichi.lang import matrix_ops 556 557 return matrix_ops.diag(dim, val)
Returns a diagonal square matrix with the diagonals filled
with val.
Args: dim (int): the dimension of the wanted square matrix. val (TypeVar): value for the diagonal elements.
Returns:
~taichi.Matrix: The wanted diagonal matrix.
Example::
>>> m = ti.Matrix.diag(3, 1)
[[1, 0, 0],
[0, 1, 0],
[0, 0, 1]]
559 def sum(self): 560 """Return the sum of all elements. 561 562 Example:: 563 564 >>> m = ti.Matrix([[1, 2], [3, 4]]) 565 >>> m.sum() 566 10 567 """ 568 # pylint: disable=C0415 569 from taichi.lang import matrix_ops 570 571 return matrix_ops.sum(self)
Return the sum of all elements.
Example::
>>> m = ti.Matrix([[1, 2], [3, 4]])
>>> m.sum()
10
573 def norm(self, eps=0): 574 """Returns the square root of the sum of the absolute squares 575 of its elements. 576 577 Args: 578 eps (Number): a safe-guard value for sqrt, usually 0. 579 580 Example:: 581 582 >>> a = ti.Vector([3, 4]) 583 >>> a.norm() 584 5 585 586 Returns: 587 The square root of the sum of the absolute squares of its elements. 588 """ 589 # pylint: disable=C0415 590 from taichi.lang import matrix_ops 591 592 return matrix_ops.norm(self, eps=eps)
Returns the square root of the sum of the absolute squares of its elements.
Args: eps (Number): a safe-guard value for sqrt, usually 0.
Example::
>>> a = ti.Vector([3, 4])
>>> a.norm()
5
Returns: The square root of the sum of the absolute squares of its elements.
594 def norm_inv(self, eps=0): 595 """The inverse of the matrix :func:`~taichi.lang.matrix.Matrix.norm`. 596 597 Args: 598 eps (float): a safe-guard value for sqrt, usually 0. 599 600 Returns: 601 The inverse of the matrix/vector `norm`. 602 """ 603 # pylint: disable=C0415 604 from taichi.lang import matrix_ops 605 606 return matrix_ops.norm_inv(self, eps=eps)
The inverse of the matrix ~taichi.lang.matrix.Matrix.norm().
Args: eps (float): a safe-guard value for sqrt, usually 0.
Returns:
The inverse of the matrix/vector norm.
608 def norm_sqr(self): 609 """Returns the sum of the absolute squares of its elements.""" 610 # pylint: disable=C0415 611 from taichi.lang import matrix_ops 612 613 return matrix_ops.norm_sqr(self)
Returns the sum of the absolute squares of its elements.
615 def max(self): 616 """Returns the maximum element value.""" 617 # pylint: disable=C0415 618 from taichi.lang import matrix_ops 619 620 return matrix_ops.max(self)
Returns the maximum element value.
622 def min(self): 623 """Returns the minimum element value.""" 624 # pylint: disable=C0415 625 from taichi.lang import matrix_ops 626 627 return matrix_ops.min(self)
Returns the minimum element value.
629 def any(self): 630 """Test whether any element not equal zero. 631 632 Returns: 633 bool: `True` if any element is not equal zero, `False` otherwise. 634 635 Example:: 636 637 >>> v = ti.Vector([0, 0, 1]) 638 >>> v.any() 639 True 640 """ 641 # pylint: disable=C0415 642 from taichi.lang import matrix_ops 643 644 return matrix_ops.any(self)
Test whether any element not equal zero.
Returns:
bool: True if any element is not equal zero, False otherwise.
Example::
>>> v = ti.Vector([0, 0, 1])
>>> v.any()
True
646 def all(self): 647 """Test whether all element not equal zero. 648 649 Returns: 650 bool: `True` if all elements are not equal zero, `False` otherwise. 651 652 Example:: 653 654 >>> v = ti.Vector([0, 0, 1]) 655 >>> v.all() 656 False 657 """ 658 # pylint: disable=C0415 659 from taichi.lang import matrix_ops 660 661 return matrix_ops.all(self)
Test whether all element not equal zero.
Returns:
bool: True if all elements are not equal zero, False otherwise.
Example::
>>> v = ti.Vector([0, 0, 1])
>>> v.all()
False
663 def fill(self, val): 664 """Fills the matrix with a specified value. 665 666 Args: 667 val (Union[int, float]): Value to fill. 668 669 Example:: 670 671 >>> A = ti.Matrix([0, 1, 2, 3]) 672 >>> A.fill(-1) 673 >>> A 674 [-1, -1, -1, -1] 675 """ 676 # pylint: disable=C0415 677 from taichi.lang import matrix_ops 678 679 return matrix_ops.fill(self, val)
Fills the matrix with a specified value.
Args: val (Union[int, float]): Value to fill.
Example::
>>> A = ti.Matrix([0, 1, 2, 3])
>>> A.fill(-1)
>>> A
[-1, -1, -1, -1]
681 def to_numpy(self): 682 """Converts this matrix to a numpy array. 683 684 Returns: 685 numpy.ndarray: The result numpy array. 686 687 Example:: 688 689 >>> A = ti.Matrix([[0], [1], [2], [3]]) 690 >>> A.to_numpy() 691 >>> A 692 array([[0], [1], [2], [3]]) 693 """ 694 if self.is_host_access: 695 return np.array(self.to_list()) 696 return self.entries
Converts this matrix to a numpy array.
Returns: numpy.ndarray: The result numpy array.
Example::
>>> A = ti.Matrix([[0], [1], [2], [3]])
>>> A.to_numpy()
>>> A
array([[0], [1], [2], [3]])
734 @staticmethod 735 @taichi_scope 736 def zero(dt, n, m=None): 737 """Constructs a Matrix filled with zeros. 738 739 Args: 740 dt (DataType): The desired data type. 741 n (int): The first dimension (row) of the matrix. 742 m (int, optional): The second dimension (column) of the matrix. 743 744 Returns: 745 :class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance filled with zeros. 746 747 """ 748 from taichi.lang import matrix_ops # pylint: disable=C0415 749 750 if m is None: 751 return matrix_ops._filled_vector(n, dt, 0) 752 return matrix_ops._filled_matrix(n, m, dt, 0)
Constructs a Matrix filled with zeros.
Args: dt (DataType): The desired data type. n (int): The first dimension (row) of the matrix. m (int, optional): The second dimension (column) of the matrix.
Returns:
~taichi.lang.matrix.Matrix: A ~taichi.lang.matrix.Matrix instance filled with zeros.
754 @staticmethod 755 @taichi_scope 756 def one(dt, n, m=None): 757 """Constructs a Matrix filled with ones. 758 759 Args: 760 dt (DataType): The desired data type. 761 n (int): The first dimension (row) of the matrix. 762 m (int, optional): The second dimension (column) of the matrix. 763 764 Returns: 765 :class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance filled with ones. 766 767 """ 768 from taichi.lang import matrix_ops # pylint: disable=C0415 769 770 if m is None: 771 return matrix_ops._filled_vector(n, dt, 1) 772 return matrix_ops._filled_matrix(n, m, dt, 1)
Constructs a Matrix filled with ones.
Args: dt (DataType): The desired data type. n (int): The first dimension (row) of the matrix. m (int, optional): The second dimension (column) of the matrix.
Returns:
~taichi.lang.matrix.Matrix: A ~taichi.lang.matrix.Matrix instance filled with ones.
774 @staticmethod 775 @taichi_scope 776 def unit(n, i, dt=None): 777 """Constructs a n-D vector with the `i`-th entry being equal to one and 778 the remaining entries are all zeros. 779 780 Args: 781 n (int): The length of the vector. 782 i (int): The index of the entry that will be filled with one. 783 dt (:mod:`~taichi.types.primitive_types`, optional): The desired data type. 784 785 Returns: 786 :class:`~taichi.Matrix`: The returned vector. 787 788 Example:: 789 790 >>> A = ti.Matrix.unit(3, 1) 791 >>> A 792 [0, 1, 0] 793 """ 794 from taichi.lang import matrix_ops # pylint: disable=C0415 795 796 if dt is None: 797 dt = int 798 assert 0 <= i < n 799 return matrix_ops._unit_vector(n, i, dt)
Constructs a n-D vector with the i-th entry being equal to one and
the remaining entries are all zeros.
Args:
n (int): The length of the vector.
i (int): The index of the entry that will be filled with one.
dt (~taichi.types.primitive_types, optional): The desired data type.
Returns:
~taichi.Matrix: The returned vector.
Example::
>>> A = ti.Matrix.unit(3, 1)
>>> A
[0, 1, 0]
801 @staticmethod 802 @taichi_scope 803 def identity(dt, n): 804 """Constructs an identity Matrix with shape (n, n). 805 806 Args: 807 dt (DataType): The desired data type. 808 n (int): The number of rows/columns. 809 810 Returns: 811 :class:`~taichi.Matrix`: An `n x n` identity matrix. 812 """ 813 from taichi.lang import matrix_ops # pylint: disable=C0415 814 815 return matrix_ops._identity_matrix(n, dt)
Constructs an identity Matrix with shape (n, n).
Args: dt (DataType): The desired data type. n (int): The number of rows/columns.
Returns:
~taichi.Matrix: An n x n identity matrix.
817 @classmethod 818 @python_scope 819 def field( 820 cls, 821 n, 822 m, 823 dtype, 824 shape=None, 825 order=None, 826 name="", 827 offset=None, 828 needs_grad=False, 829 needs_dual=False, 830 layout=Layout.AOS, 831 ndim=None, 832 ): 833 """Construct a data container to hold all elements of the Matrix. 834 835 Args: 836 n (int): The desired number of rows of the Matrix. 837 m (int): The desired number of columns of the Matrix. 838 dtype (DataType, optional): The desired data type of the Matrix. 839 shape (Union[int, tuple of int], optional): The desired shape of the Matrix. 840 order (str, optional): order of the shape laid out in memory. 841 name (string, optional): The custom name of the field. 842 offset (Union[int, tuple of int], optional): The coordinate offset 843 of all elements in a field. 844 needs_grad (bool, optional): Whether the Matrix need grad field (reverse mode autodiff). 845 needs_dual (bool, optional): Whether the Matrix need dual field (forward mode autodiff). 846 layout (Layout, optional): The field layout, either Array Of 847 Structure (AOS) or Structure Of Array (SOA). 848 849 Returns: 850 :class:`~taichi.Matrix`: A matrix. 851 """ 852 entries = [] 853 element_dim = ndim if ndim is not None else 2 854 if isinstance(dtype, (list, tuple, np.ndarray)): 855 # set different dtype for each element in Matrix 856 # see #2135 857 if m == 1: 858 assert ( 859 len(np.shape(dtype)) == 1 and len(dtype) == n 860 ), f"Please set correct dtype list for Vector. The shape of dtype list should be ({n}, ) instead of {np.shape(dtype)}" 861 for i in range(n): 862 entries.append( 863 impl.create_field_member( 864 dtype[i], 865 name=name, 866 needs_grad=needs_grad, 867 needs_dual=needs_dual, 868 ) 869 ) 870 else: 871 assert ( 872 len(np.shape(dtype)) == 2 and len(dtype) == n and len(dtype[0]) == m 873 ), f"Please set correct dtype list for Matrix. The shape of dtype list should be ({n}, {m}) instead of {np.shape(dtype)}" 874 for i in range(n): 875 for j in range(m): 876 entries.append( 877 impl.create_field_member( 878 dtype[i][j], 879 name=name, 880 needs_grad=needs_grad, 881 needs_dual=needs_dual, 882 ) 883 ) 884 else: 885 for _ in range(n * m): 886 entries.append(impl.create_field_member(dtype, name=name, needs_grad=needs_grad, needs_dual=needs_dual)) 887 entries, entries_grad, entries_dual = zip(*entries) 888 889 entries = MatrixField(entries, n, m, element_dim) 890 if all(entries_grad): 891 entries_grad = MatrixField(entries_grad, n, m, element_dim) 892 entries._set_grad(entries_grad) 893 if all(entries_dual): 894 entries_dual = MatrixField(entries_dual, n, m, element_dim) 895 entries._set_dual(entries_dual) 896 897 impl.get_runtime().matrix_fields.append(entries) 898 899 if shape is None: 900 if offset is not None: 901 raise TaichiSyntaxError("shape cannot be None when offset is set") 902 if order is not None: 903 raise TaichiSyntaxError("shape cannot be None when order is set") 904 else: 905 if isinstance(shape, numbers.Number): 906 shape = (shape,) 907 if isinstance(offset, numbers.Number): 908 offset = (offset,) 909 dim = len(shape) 910 if offset is not None and dim != len(offset): 911 raise TaichiSyntaxError( 912 f"The dimensionality of shape and offset must be the same ({dim} != {len(offset)})" 913 ) 914 axis_seq = [] 915 shape_seq = [] 916 if order is not None: 917 if dim != len(order): 918 raise TaichiSyntaxError( 919 f"The dimensionality of shape and order must be the same ({dim} != {len(order)})" 920 ) 921 if dim != len(set(order)): 922 raise TaichiSyntaxError("The axes in order must be different") 923 for ch in order: 924 axis = ord(ch) - ord("i") 925 if axis < 0 or axis >= dim: 926 raise TaichiSyntaxError(f"Invalid axis {ch}") 927 axis_seq.append(axis) 928 shape_seq.append(shape[axis]) 929 else: 930 axis_seq = list(range(dim)) 931 shape_seq = list(shape) 932 same_level = order is None 933 if layout == Layout.SOA: 934 for e in entries._get_field_members(): 935 impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset) 936 if needs_grad: 937 for e in entries_grad._get_field_members(): 938 impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset) 939 if needs_dual: 940 for e in entries_dual._get_field_members(): 941 impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset) 942 else: 943 impl._create_snode(axis_seq, shape_seq, same_level).place(entries, offset=offset) 944 if needs_grad: 945 impl._create_snode(axis_seq, shape_seq, same_level).place(entries_grad, offset=offset) 946 if needs_dual: 947 impl._create_snode(axis_seq, shape_seq, same_level).place(entries_dual, offset=offset) 948 return entries
Construct a data container to hold all elements of the Matrix.
Args: n (int): The desired number of rows of the Matrix. m (int): The desired number of columns of the Matrix. dtype (DataType, optional): The desired data type of the Matrix. shape (Union[int, tuple of int], optional): The desired shape of the Matrix. order (str, optional): order of the shape laid out in memory. name (string, optional): The custom name of the field. offset (Union[int, tuple of int], optional): The coordinate offset of all elements in a field. needs_grad (bool, optional): Whether the Matrix need grad field (reverse mode autodiff). needs_dual (bool, optional): Whether the Matrix need dual field (forward mode autodiff). layout (Layout, optional): The field layout, either Array Of Structure (AOS) or Structure Of Array (SOA).
Returns:
~taichi.Matrix: A matrix.
950 @classmethod 951 @python_scope 952 def ndarray(cls, n, m, dtype, shape): 953 """Defines a Taichi ndarray with matrix elements. 954 This function must be called in Python scope, and after `ti.init` is called. 955 956 Args: 957 n (int): Number of rows of the matrix. 958 m (int): Number of columns of the matrix. 959 dtype (DataType): Data type of each value. 960 shape (Union[int, tuple[int]]): Shape of the ndarray. 961 962 Example:: 963 964 The code below shows how a Taichi ndarray with matrix elements \ 965 can be declared and defined:: 966 967 >>> x = ti.Matrix.ndarray(4, 5, ti.f32, shape=(16, 8)) 968 """ 969 if isinstance(shape, numbers.Number): 970 shape = (shape,) 971 return MatrixNdarray(n, m, dtype, shape)
Defines a Taichi ndarray with matrix elements.
This function must be called in Python scope, and after ti.init is called.
Args: n (int): Number of rows of the matrix. m (int): Number of columns of the matrix. dtype (DataType): Data type of each value. shape (Union[int, tuple[int]]): Shape of the ndarray.
Example::
The code below shows how a Taichi ndarray with matrix elements can be declared and defined::
>>> x = ti.Matrix.ndarray(4, 5, ti.f32, shape=(16, 8))
973 @staticmethod 974 def rows(rows): 975 """Constructs a matrix by concatenating a list of 976 vectors/lists row by row. Must be called in Taichi scope. 977 978 Args: 979 rows (List): A list of Vector (1-D Matrix) or a list of list. 980 981 Returns: 982 :class:`~taichi.Matrix`: A matrix. 983 984 Example:: 985 986 >>> @ti.kernel 987 >>> def test(): 988 >>> v1 = ti.Vector([1, 2, 3]) 989 >>> v2 = ti.Vector([4, 5, 6]) 990 >>> m = ti.Matrix.rows([v1, v2]) 991 >>> print(m) 992 >>> 993 >>> test() 994 [[1, 2, 3], [4, 5, 6]] 995 """ 996 from taichi.lang import matrix_ops # pylint: disable=C0415 997 998 return matrix_ops.rows(rows)
Constructs a matrix by concatenating a list of vectors/lists row by row. Must be called in Taichi scope.
Args: rows (List): A list of Vector (1-D Matrix) or a list of list.
Returns:
~taichi.Matrix: A matrix.
Example::
>>> @ti.kernel
>>> def test():
>>> v1 = ti.Vector([1, 2, 3])
>>> v2 = ti.Vector([4, 5, 6])
>>> m = ti.Matrix.rows([v1, v2])
>>> print(m)
>>>
>>> test()
[[1, 2, 3], [4, 5, 6]]
1000 @staticmethod 1001 def cols(cols): 1002 """Constructs a Matrix instance by concatenating Vectors/lists column by column. 1003 1004 Args: 1005 cols (List): A list of Vector (1-D Matrix) or a list of list. 1006 1007 Returns: 1008 :class:`~taichi.Matrix`: A matrix. 1009 1010 Example:: 1011 1012 >>> @ti.kernel 1013 >>> def test(): 1014 >>> v1 = ti.Vector([1, 2, 3]) 1015 >>> v2 = ti.Vector([4, 5, 6]) 1016 >>> m = ti.Matrix.cols([v1, v2]) 1017 >>> print(m) 1018 >>> 1019 >>> test() 1020 [[1, 4], [2, 5], [3, 6]] 1021 """ 1022 from taichi.lang import matrix_ops # pylint: disable=C0415 1023 1024 return matrix_ops.cols(cols)
Constructs a Matrix instance by concatenating Vectors/lists column by column.
Args: cols (List): A list of Vector (1-D Matrix) or a list of list.
Returns:
~taichi.Matrix: A matrix.
Example::
>>> @ti.kernel
>>> def test():
>>> v1 = ti.Vector([1, 2, 3])
>>> v2 = ti.Vector([4, 5, 6])
>>> m = ti.Matrix.cols([v1, v2])
>>> print(m)
>>>
>>> test()
[[1, 4], [2, 5], [3, 6]]
1032 def dot(self, other): 1033 """Performs the dot product of two vectors. 1034 1035 To call this method, both multiplicatives must be vectors. 1036 1037 Args: 1038 other (:class:`~taichi.Matrix`): The input Vector. 1039 1040 Returns: 1041 DataType: The dot product result (scalar) of the two Vectors. 1042 1043 Example:: 1044 1045 >>> v1 = ti.Vector([1, 2, 3]) 1046 >>> v2 = ti.Vector([3, 4, 5]) 1047 >>> v1.dot(v2) 1048 26 1049 """ 1050 from taichi.lang import matrix_ops # pylint: disable=C0415 1051 1052 return matrix_ops.dot(self, other)
Performs the dot product of two vectors.
To call this method, both multiplicatives must be vectors.
Args:
other (~taichi.Matrix): The input Vector.
Returns: DataType: The dot product result (scalar) of the two Vectors.
Example::
>>> v1 = ti.Vector([1, 2, 3])
>>> v2 = ti.Vector([3, 4, 5])
>>> v1.dot(v2)
26
1054 def cross(self, other): 1055 """Performs the cross product with the input vector (1-D Matrix). 1056 1057 Both two vectors must have the same dimension <= 3. 1058 1059 For two 2d vectors (x1, y1) and (x2, y2), the return value is the 1060 scalar `x1*y2 - x2*y1`. 1061 1062 For two 3d vectors `v` and `w`, the return value is the 3d vector 1063 `v x w`. 1064 1065 Args: 1066 other (:class:`~taichi.Matrix`): The input Vector. 1067 1068 Returns: 1069 :class:`~taichi.Matrix`: The cross product of the two Vectors. 1070 """ 1071 from taichi.lang import matrix_ops # pylint: disable=C0415 1072 1073 return matrix_ops.cross(self, other)
Performs the cross product with the input vector (1-D Matrix).
Both two vectors must have the same dimension <= 3.
For two 2d vectors (x1, y1) and (x2, y2), the return value is the
scalar x1*y2 - x2*y1.
For two 3d vectors v and w, the return value is the 3d vector
v x w.
Args:
other (~taichi.Matrix): The input Vector.
Returns:
~taichi.Matrix: The cross product of the two Vectors.
1075 def outer_product(self, other): 1076 """Performs the outer product with the input Vector (1-D Matrix). 1077 1078 The outer_product of two vectors `v = (x1, x2, ..., xn)`, 1079 `w = (y1, y2, ..., yn)` is a `n` times `n` square matrix, and its `(i, j)` 1080 entry is equal to `xi*yj`. 1081 1082 Args: 1083 other (:class:`~taichi.Matrix`): The input Vector. 1084 1085 Returns: 1086 :class:`~taichi.Matrix`: The outer product of the two Vectors. 1087 """ 1088 from taichi.lang import matrix_ops # pylint: disable=C0415 1089 1090 return matrix_ops.outer_product(self, other)
Performs the outer product with the input Vector (1-D Matrix).
The outer_product of two vectors v = (x1, x2, ..., xn),
w = (y1, y2, ..., yn) is a n times n square matrix, and its (i, j)
entry is equal to xi*yj.
Args:
other (~taichi.Matrix): The input Vector.
Returns:
~taichi.Matrix: The outer product of the two Vectors.
1146class MatrixField(Field): 1147 """Taichi matrix field with SNode implementation. 1148 1149 Args: 1150 vars (List[Expr]): Field members. 1151 n (Int): Number of rows. 1152 m (Int): Number of columns. 1153 ndim (Int): Number of dimensions; forced reshape if given. 1154 """ 1155 1156 def __init__(self, _vars, n, m, ndim=2): 1157 assert len(_vars) == n * m 1158 assert ndim in (0, 1, 2) 1159 super().__init__(_vars) 1160 self.n = n 1161 self.m = m 1162 self.ndim = ndim 1163 self.ptr = ti_python_core.expr_matrix_field([var.ptr for var in self.vars], [n, m][:ndim]) 1164 1165 def get_scalar_field(self, *indices): 1166 """Creates a ScalarField using a specific field member. 1167 1168 Args: 1169 indices (Tuple[Int]): Specified indices of the field member. 1170 1171 Returns: 1172 ScalarField: The result ScalarField. 1173 """ 1174 assert len(indices) in [1, 2] 1175 i = indices[0] 1176 j = 0 if len(indices) == 1 else indices[1] 1177 return ScalarField(self.vars[i * self.m + j]) 1178 1179 def _get_dynamic_index_stride(self): 1180 if self.ptr.get_dynamic_indexable(): 1181 return self.ptr.get_dynamic_index_stride() 1182 return None 1183 1184 def _calc_dynamic_index_stride(self): 1185 # Algorithm: https://github.com/taichi-dev/taichi/issues/3810 1186 paths = [ScalarField(var).snode._path_from_root() for var in self.vars] 1187 num_members = len(paths) 1188 if num_members == 1: 1189 self.ptr.set_dynamic_index_stride(0) 1190 return 1191 length = len(paths[0]) 1192 if any(len(path) != length or ti_python_core.is_quant(path[length - 1]._dtype) for path in paths): 1193 return 1194 for i in range(length): 1195 if any(path[i] != paths[0][i] for path in paths): 1196 depth_below_lca = i 1197 break 1198 for i in range(depth_below_lca, length - 1): 1199 if any( 1200 path[i].ptr.type != ti_python_core.SNodeType.dense 1201 or path[i]._cell_size_bytes != paths[0][i]._cell_size_bytes 1202 or path[i + 1]._offset_bytes_in_parent_cell != paths[0][i + 1]._offset_bytes_in_parent_cell 1203 for path in paths 1204 ): 1205 return 1206 stride = ( 1207 paths[1][depth_below_lca]._offset_bytes_in_parent_cell 1208 - paths[0][depth_below_lca]._offset_bytes_in_parent_cell 1209 ) 1210 for i in range(2, num_members): 1211 if ( 1212 stride 1213 != paths[i][depth_below_lca]._offset_bytes_in_parent_cell 1214 - paths[i - 1][depth_below_lca]._offset_bytes_in_parent_cell 1215 ): 1216 return 1217 self.ptr.set_dynamic_index_stride(stride) 1218 1219 def fill(self, val): 1220 """Fills this matrix field with specified values. 1221 1222 Args: 1223 val (Union[Number, Expr, List, Tuple, Matrix]): Values to fill, 1224 should have consistent dimension consistent with `self`. 1225 """ 1226 if isinstance(val, numbers.Number) or (isinstance(val, expr.Expr) and not val.is_tensor()): 1227 if self.ndim == 2: 1228 val = tuple(tuple(val for _ in range(self.m)) for _ in range(self.n)) 1229 else: 1230 assert self.ndim == 1 1231 val = tuple(val for _ in range(self.n)) 1232 elif isinstance(val, expr.Expr) and val.is_tensor(): 1233 assert val.n == self.n 1234 if self.ndim != 1: 1235 assert val.m == self.m 1236 else: 1237 if isinstance(val, Matrix): 1238 val = val.to_list() 1239 assert isinstance(val, (list, tuple)) 1240 val = tuple(tuple(x) if isinstance(x, list) else x for x in val) 1241 assert len(val) == self.n 1242 if self.ndim != 1: 1243 assert len(val[0]) == self.m 1244 if in_python_scope(): 1245 from taichi._kernels import field_fill_python_scope # pylint: disable=C0415 1246 1247 field_fill_python_scope(self, val) 1248 else: 1249 from taichi._funcs import field_fill_taichi_scope # pylint: disable=C0415 1250 1251 field_fill_taichi_scope(self, val) 1252 1253 @python_scope 1254 def to_numpy(self, keep_dims=False, dtype=None): 1255 """Converts the field instance to a NumPy array. 1256 1257 Args: 1258 keep_dims (bool, optional): Whether to keep the dimension after conversion. 1259 When keep_dims=True, on an n-D matrix field, the numpy array always has n+2 dims, even for 1x1, 1xn, nx1 matrix fields. 1260 When keep_dims=False, the resulting numpy array should skip the matrix dims with size 1. 1261 For example, a 4x1 or 1x4 matrix field with 5x6x7 elements results in an array of shape 5x6x7x4. 1262 dtype (DataType, optional): The desired data type of returned numpy array. 1263 1264 Returns: 1265 numpy.ndarray: The result NumPy array. 1266 """ 1267 if dtype is None: 1268 dtype = to_numpy_type(self.dtype) 1269 as_vector = self.m == 1 and not keep_dims 1270 shape_ext = (self.n,) if as_vector else (self.n, self.m) 1271 arr = np.zeros(self.shape + shape_ext, dtype=dtype) 1272 from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415 1273 1274 matrix_to_ext_arr(self, arr, as_vector) 1275 runtime_ops.sync() 1276 return arr 1277 1278 def to_torch(self, device=None, keep_dims=False): 1279 """Converts the field instance to a PyTorch tensor. 1280 1281 Args: 1282 device (torch.device, optional): The desired device of returned tensor. 1283 keep_dims (bool, optional): Whether to keep the dimension after conversion. 1284 See :meth:`~taichi.lang.field.MatrixField.to_numpy` for more detailed explanation. 1285 1286 Returns: 1287 torch.tensor: The result torch tensor. 1288 """ 1289 import torch # pylint: disable=C0415 1290 1291 as_vector = self.m == 1 and not keep_dims 1292 shape_ext = (self.n,) if as_vector else (self.n, self.m) 1293 # pylint: disable=E1101 1294 arr = torch.empty(self.shape + shape_ext, dtype=to_pytorch_type(self.dtype), device=device) 1295 from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415 1296 1297 matrix_to_ext_arr(self, arr, as_vector) 1298 runtime_ops.sync() 1299 return arr 1300 1301 def to_paddle(self, place=None, keep_dims=False): 1302 """Converts the field instance to a Paddle tensor. 1303 1304 Args: 1305 place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor. 1306 keep_dims (bool, optional): Whether to keep the dimension after conversion. 1307 See :meth:`~taichi.lang.field.MatrixField.to_numpy` for more detailed explanation. 1308 1309 Returns: 1310 paddle.Tensor: The result paddle tensor. 1311 """ 1312 import paddle # pylint: disable=C0415 1313 1314 as_vector = self.m == 1 and not keep_dims and self.ndim == 1 1315 shape_ext = (self.n,) if as_vector else (self.n, self.m) 1316 # pylint: disable=E1101 1317 # paddle.empty() doesn't support argument `place`` 1318 arr = paddle.to_tensor( 1319 paddle.empty(self.shape + shape_ext, to_paddle_type(self.dtype)), 1320 place=place, 1321 ) 1322 from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415 1323 1324 matrix_to_ext_arr(self, arr, as_vector) 1325 runtime_ops.sync() 1326 return arr 1327 1328 @python_scope 1329 def _from_external_arr(self, arr): 1330 if len(arr.shape) == len(self.shape) + 1: 1331 as_vector = True 1332 assert self.m == 1, "This is not a vector field" 1333 else: 1334 as_vector = False 1335 assert len(arr.shape) == len(self.shape) + 2 1336 dim_ext = 1 if as_vector else 2 1337 assert len(arr.shape) == len(self.shape) + dim_ext 1338 from taichi._kernels import ext_arr_to_matrix # pylint: disable=C0415 1339 1340 ext_arr_to_matrix(arr, self, as_vector) 1341 runtime_ops.sync() 1342 1343 @python_scope 1344 def from_numpy(self, arr): 1345 """Copies an `numpy.ndarray` into this field. 1346 1347 Example:: 1348 1349 >>> m = ti.Matrix.field(2, 2, ti.f32, shape=(3, 3)) 1350 >>> arr = numpy.ones((3, 3, 2, 2)) 1351 >>> m.from_numpy(arr) 1352 """ 1353 1354 if not arr.flags.c_contiguous: 1355 arr = np.ascontiguousarray(arr) 1356 self._from_external_arr(arr) 1357 1358 @python_scope 1359 def __setitem__(self, key, value): 1360 self._initialize_host_accessors() 1361 self[key]._set_entries(value) 1362 1363 @python_scope 1364 def __getitem__(self, key): 1365 self._initialize_host_accessors() 1366 key = self._pad_key(key) 1367 _host_access = self._host_access(key) 1368 if self.ndim == 1: 1369 return Vector([_host_access[i] for i in range(self.n)]) 1370 return Matrix([[_host_access[i * self.m + j] for j in range(self.m)] for i in range(self.n)]) 1371 1372 def __repr__(self): 1373 # make interactive shell happy, prevent materialization 1374 return f"<{self.n}x{self.m} ti.Matrix.field>"
Taichi matrix field with SNode implementation.
Args: vars (List[Expr]): Field members. n (Int): Number of rows. m (Int): Number of columns. ndim (Int): Number of dimensions; forced reshape if given.
1165 def get_scalar_field(self, *indices): 1166 """Creates a ScalarField using a specific field member. 1167 1168 Args: 1169 indices (Tuple[Int]): Specified indices of the field member. 1170 1171 Returns: 1172 ScalarField: The result ScalarField. 1173 """ 1174 assert len(indices) in [1, 2] 1175 i = indices[0] 1176 j = 0 if len(indices) == 1 else indices[1] 1177 return ScalarField(self.vars[i * self.m + j])
Creates a ScalarField using a specific field member.
Args: indices (Tuple[Int]): Specified indices of the field member.
Returns: ScalarField: The result ScalarField.
1219 def fill(self, val): 1220 """Fills this matrix field with specified values. 1221 1222 Args: 1223 val (Union[Number, Expr, List, Tuple, Matrix]): Values to fill, 1224 should have consistent dimension consistent with `self`. 1225 """ 1226 if isinstance(val, numbers.Number) or (isinstance(val, expr.Expr) and not val.is_tensor()): 1227 if self.ndim == 2: 1228 val = tuple(tuple(val for _ in range(self.m)) for _ in range(self.n)) 1229 else: 1230 assert self.ndim == 1 1231 val = tuple(val for _ in range(self.n)) 1232 elif isinstance(val, expr.Expr) and val.is_tensor(): 1233 assert val.n == self.n 1234 if self.ndim != 1: 1235 assert val.m == self.m 1236 else: 1237 if isinstance(val, Matrix): 1238 val = val.to_list() 1239 assert isinstance(val, (list, tuple)) 1240 val = tuple(tuple(x) if isinstance(x, list) else x for x in val) 1241 assert len(val) == self.n 1242 if self.ndim != 1: 1243 assert len(val[0]) == self.m 1244 if in_python_scope(): 1245 from taichi._kernels import field_fill_python_scope # pylint: disable=C0415 1246 1247 field_fill_python_scope(self, val) 1248 else: 1249 from taichi._funcs import field_fill_taichi_scope # pylint: disable=C0415 1250 1251 field_fill_taichi_scope(self, val)
Fills this matrix field with specified values.
Args:
val (Union[Number, Expr, List, Tuple, Matrix]): Values to fill,
should have consistent dimension consistent with self.
1253 @python_scope 1254 def to_numpy(self, keep_dims=False, dtype=None): 1255 """Converts the field instance to a NumPy array. 1256 1257 Args: 1258 keep_dims (bool, optional): Whether to keep the dimension after conversion. 1259 When keep_dims=True, on an n-D matrix field, the numpy array always has n+2 dims, even for 1x1, 1xn, nx1 matrix fields. 1260 When keep_dims=False, the resulting numpy array should skip the matrix dims with size 1. 1261 For example, a 4x1 or 1x4 matrix field with 5x6x7 elements results in an array of shape 5x6x7x4. 1262 dtype (DataType, optional): The desired data type of returned numpy array. 1263 1264 Returns: 1265 numpy.ndarray: The result NumPy array. 1266 """ 1267 if dtype is None: 1268 dtype = to_numpy_type(self.dtype) 1269 as_vector = self.m == 1 and not keep_dims 1270 shape_ext = (self.n,) if as_vector else (self.n, self.m) 1271 arr = np.zeros(self.shape + shape_ext, dtype=dtype) 1272 from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415 1273 1274 matrix_to_ext_arr(self, arr, as_vector) 1275 runtime_ops.sync() 1276 return arr
Converts the field instance to a NumPy array.
Args: keep_dims (bool, optional): Whether to keep the dimension after conversion. When keep_dims=True, on an n-D matrix field, the numpy array always has n+2 dims, even for 1x1, 1xn, nx1 matrix fields. When keep_dims=False, the resulting numpy array should skip the matrix dims with size 1. For example, a 4x1 or 1x4 matrix field with 5x6x7 elements results in an array of shape 5x6x7x4. dtype (DataType, optional): The desired data type of returned numpy array.
Returns: numpy.ndarray: The result NumPy array.
1278 def to_torch(self, device=None, keep_dims=False): 1279 """Converts the field instance to a PyTorch tensor. 1280 1281 Args: 1282 device (torch.device, optional): The desired device of returned tensor. 1283 keep_dims (bool, optional): Whether to keep the dimension after conversion. 1284 See :meth:`~taichi.lang.field.MatrixField.to_numpy` for more detailed explanation. 1285 1286 Returns: 1287 torch.tensor: The result torch tensor. 1288 """ 1289 import torch # pylint: disable=C0415 1290 1291 as_vector = self.m == 1 and not keep_dims 1292 shape_ext = (self.n,) if as_vector else (self.n, self.m) 1293 # pylint: disable=E1101 1294 arr = torch.empty(self.shape + shape_ext, dtype=to_pytorch_type(self.dtype), device=device) 1295 from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415 1296 1297 matrix_to_ext_arr(self, arr, as_vector) 1298 runtime_ops.sync() 1299 return arr
Converts the field instance to a PyTorch tensor.
Args:
device (torch.device, optional): The desired device of returned tensor.
keep_dims (bool, optional): Whether to keep the dimension after conversion.
See ~taichi.lang.field.MatrixField.to_numpy() for more detailed explanation.
Returns: torch.tensor: The result torch tensor.
1301 def to_paddle(self, place=None, keep_dims=False): 1302 """Converts the field instance to a Paddle tensor. 1303 1304 Args: 1305 place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor. 1306 keep_dims (bool, optional): Whether to keep the dimension after conversion. 1307 See :meth:`~taichi.lang.field.MatrixField.to_numpy` for more detailed explanation. 1308 1309 Returns: 1310 paddle.Tensor: The result paddle tensor. 1311 """ 1312 import paddle # pylint: disable=C0415 1313 1314 as_vector = self.m == 1 and not keep_dims and self.ndim == 1 1315 shape_ext = (self.n,) if as_vector else (self.n, self.m) 1316 # pylint: disable=E1101 1317 # paddle.empty() doesn't support argument `place`` 1318 arr = paddle.to_tensor( 1319 paddle.empty(self.shape + shape_ext, to_paddle_type(self.dtype)), 1320 place=place, 1321 ) 1322 from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415 1323 1324 matrix_to_ext_arr(self, arr, as_vector) 1325 runtime_ops.sync() 1326 return arr
Converts the field instance to a Paddle tensor.
Args:
place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor.
keep_dims (bool, optional): Whether to keep the dimension after conversion.
See ~taichi.lang.field.MatrixField.to_numpy() for more detailed explanation.
Returns: paddle.Tensor: The result paddle tensor.
1343 @python_scope 1344 def from_numpy(self, arr): 1345 """Copies an `numpy.ndarray` into this field. 1346 1347 Example:: 1348 1349 >>> m = ti.Matrix.field(2, 2, ti.f32, shape=(3, 3)) 1350 >>> arr = numpy.ones((3, 3, 2, 2)) 1351 >>> m.from_numpy(arr) 1352 """ 1353 1354 if not arr.flags.c_contiguous: 1355 arr = np.ascontiguousarray(arr) 1356 self._from_external_arr(arr)
Copies an numpy.ndarray into this field.
Example::
>>> m = ti.Matrix.field(2, 2, ti.f32, shape=(3, 3))
>>> arr = numpy.ones((3, 3, 2, 2))
>>> m.from_numpy(arr)
Inherited Members
1652class MatrixNdarray(Ndarray): 1653 """Taichi ndarray with matrix elements. 1654 1655 Args: 1656 n (int): Number of rows of the matrix. 1657 m (int): Number of columns of the matrix. 1658 dtype (DataType): Data type of each value. 1659 shape (Union[int, tuple[int]]): Shape of the ndarray. 1660 1661 Example:: 1662 1663 >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3)) 1664 """ 1665 1666 def __init__(self, n, m, dtype, shape): 1667 self.n = n 1668 self.m = m 1669 super().__init__() 1670 # TODO(zhanlue): remove self.dtype and migrate its usages to element_type 1671 self.dtype = cook_dtype(dtype) 1672 1673 self.layout = Layout.AOS 1674 self.shape = tuple(shape) 1675 self.element_type = _type_factory.get_tensor_type((self.n, self.m), self.dtype) 1676 # TODO: we should pass in element_type, shape, layout instead. 1677 self.arr = impl.get_runtime().prog.create_ndarray( 1678 cook_dtype(self.element_type), 1679 shape, 1680 Layout.AOS, 1681 zero_fill=True, 1682 dbg_info=ti_python_core.DebugInfo(get_traceback()), 1683 ) 1684 1685 @property 1686 def element_shape(self): 1687 """Returns the shape of each element (a 2D matrix) in this ndarray. 1688 1689 Example:: 1690 1691 >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3)) 1692 >>> arr.element_shape 1693 (2, 2) 1694 """ 1695 return tuple(self.arr.element_shape()) 1696 1697 @python_scope 1698 def __setitem__(self, key, value): 1699 if not isinstance(value, (list, tuple)): 1700 value = list(value) 1701 if not isinstance(value[0], (list, tuple)): 1702 value = [[i] for i in value] 1703 for i in range(self.n): 1704 for j in range(self.m): 1705 self[key][i, j] = value[i][j] 1706 1707 @python_scope 1708 def __getitem__(self, key): 1709 key = () if key is None else (key,) if isinstance(key, numbers.Number) else tuple(key) 1710 return Matrix([[NdarrayHostAccess(self, key, (i, j)) for j in range(self.m)] for i in range(self.n)]) 1711 1712 @python_scope 1713 def to_numpy(self): 1714 """Converts this ndarray to a `numpy.ndarray`. 1715 1716 Example:: 1717 1718 >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1)) 1719 >>> arr.to_numpy() 1720 [[[[0. 0.] 1721 [0. 0.]]] 1722 1723 [[[0. 0.] 1724 [0. 0.]]]] 1725 """ 1726 return self._ndarray_matrix_to_numpy(as_vector=0) 1727 1728 @python_scope 1729 def from_numpy(self, arr): 1730 """Copies the data of a `numpy.ndarray` into this array. 1731 1732 Example:: 1733 1734 >>> m = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1), layout=0) 1735 >>> arr = np.ones((2, 1, 2, 2)) 1736 >>> m.from_numpy(arr) 1737 """ 1738 self._ndarray_matrix_from_numpy(arr, as_vector=0) 1739 1740 @python_scope 1741 def __deepcopy__(self, memo=None): 1742 ret_arr = MatrixNdarray(self.n, self.m, self.dtype, self.shape) 1743 ret_arr.copy_from(self) 1744 return ret_arr 1745 1746 @python_scope 1747 def _fill_by_kernel(self, val): 1748 from taichi._kernels import fill_ndarray_matrix # pylint: disable=C0415 1749 1750 shape = self.element_type.shape() 1751 n = shape[0] 1752 m = 1 1753 if len(shape) > 1: 1754 m = shape[1] 1755 1756 prim_dtype = self.element_type.element_type() 1757 matrix_type = MatrixType(n, m, len(shape), prim_dtype) 1758 if isinstance(val, Matrix): 1759 value = val 1760 else: 1761 value = matrix_type(val) 1762 fill_ndarray_matrix(self, value) 1763 1764 @python_scope 1765 def __repr__(self): 1766 return f"<{self.n}x{self.m} {Layout.AOS} ti.Matrix.ndarray>"
Taichi ndarray with matrix elements.
Args: n (int): Number of rows of the matrix. m (int): Number of columns of the matrix. dtype (DataType): Data type of each value. shape (Union[int, tuple[int]]): Shape of the ndarray.
Example::
>>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3))
1666 def __init__(self, n, m, dtype, shape): 1667 self.n = n 1668 self.m = m 1669 super().__init__() 1670 # TODO(zhanlue): remove self.dtype and migrate its usages to element_type 1671 self.dtype = cook_dtype(dtype) 1672 1673 self.layout = Layout.AOS 1674 self.shape = tuple(shape) 1675 self.element_type = _type_factory.get_tensor_type((self.n, self.m), self.dtype) 1676 # TODO: we should pass in element_type, shape, layout instead. 1677 self.arr = impl.get_runtime().prog.create_ndarray( 1678 cook_dtype(self.element_type), 1679 shape, 1680 Layout.AOS, 1681 zero_fill=True, 1682 dbg_info=ti_python_core.DebugInfo(get_traceback()), 1683 )
1685 @property 1686 def element_shape(self): 1687 """Returns the shape of each element (a 2D matrix) in this ndarray. 1688 1689 Example:: 1690 1691 >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3)) 1692 >>> arr.element_shape 1693 (2, 2) 1694 """ 1695 return tuple(self.arr.element_shape())
Returns the shape of each element (a 2D matrix) in this ndarray.
Example::
>>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3))
>>> arr.element_shape
(2, 2)
1712 @python_scope 1713 def to_numpy(self): 1714 """Converts this ndarray to a `numpy.ndarray`. 1715 1716 Example:: 1717 1718 >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1)) 1719 >>> arr.to_numpy() 1720 [[[[0. 0.] 1721 [0. 0.]]] 1722 1723 [[[0. 0.] 1724 [0. 0.]]]] 1725 """ 1726 return self._ndarray_matrix_to_numpy(as_vector=0)
Converts this ndarray to a numpy.ndarray.
Example::
>>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1))
>>> arr.to_numpy()
[[[[0. 0.]
[0. 0.]]]
[[[0. 0.]
[0. 0.]]]]
1728 @python_scope 1729 def from_numpy(self, arr): 1730 """Copies the data of a `numpy.ndarray` into this array. 1731 1732 Example:: 1733 1734 >>> m = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1), layout=0) 1735 >>> arr = np.ones((2, 1, 2, 2)) 1736 >>> m.from_numpy(arr) 1737 """ 1738 self._ndarray_matrix_from_numpy(arr, as_vector=0)
Copies the data of a numpy.ndarray into this array.
Example::
>>> m = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1), layout=0)
>>> arr = np.ones((2, 1, 2, 2))
>>> m.from_numpy(arr)
502class Mesh: 503 """The Mesh type class. 504 505 MeshTaichi offers first-class support for triangular/tetrahedral meshes 506 and allows efficient computation on these irregular data structures, 507 only available for backends supporting `ti.extension.mesh`. 508 509 See more details in https://github.com/taichi-dev/meshtaichi 510 """ 511 512 def __init__(self): 513 pass 514 515 @staticmethod 516 def _create_instance(metadata: MeshMetadata) -> MeshInstance: 517 instance = MeshInstance() 518 instance.fields = {} 519 520 instance.set_num_patches(metadata.num_patches) 521 522 for element in metadata.element_fields: 523 _ti_core.set_num_elements(instance.mesh_ptr, element, metadata.num_elements[element]) 524 instance.set_patch_max_element_num(element, metadata.max_num_per_patch[element]) 525 526 element_name = element_type_name(element) 527 setattr( 528 instance, 529 element_name, 530 MeshElementField(instance, element, {}, {}, metadata.element_fields[element]["g2r"]), 531 ) 532 instance.fields[element] = getattr(instance, element_name) 533 534 instance.set_owned_offset(element, metadata.element_fields[element]["owned"]) 535 instance.set_total_offset(element, metadata.element_fields[element]["total"]) 536 instance.set_index_mapping(element, ConvType.l2g, metadata.element_fields[element]["l2g"]) 537 instance.set_index_mapping(element, ConvType.l2r, metadata.element_fields[element]["l2r"]) 538 instance.set_index_mapping(element, ConvType.g2r, metadata.element_fields[element]["g2r"]) 539 540 for rel_type in metadata.relation_fields: 541 from_order = metadata.relation_fields[rel_type]["from_order"] 542 to_order = metadata.relation_fields[rel_type]["to_order"] 543 if from_order <= to_order: 544 instance.set_relation_dynamic( 545 rel_type, 546 metadata.relation_fields[rel_type]["value"], 547 metadata.relation_fields[rel_type]["patch_offset"], 548 metadata.relation_fields[rel_type]["offset"], 549 ) 550 else: 551 instance.set_relation_fixed(rel_type, metadata.relation_fields[rel_type]["value"]) 552 553 instance._vert_position = metadata.attrs["x"] 554 instance.patcher = metadata.patcher 555 556 return instance 557 558 @staticmethod 559 def load_meta(filename): 560 with open(filename, "r") as fi: 561 data = json.loads(fi.read()) 562 return MeshMetadata(data) 563 564 @staticmethod 565 def generate_meta(data): 566 return MeshMetadata(data)
The Mesh type class.
MeshTaichi offers first-class support for triangular/tetrahedral meshes
and allows efficient computation on these irregular data structures,
only available for backends supporting ti.extension.mesh.
See more details in https://github.com/taichi-dev/meshtaichi
305class MeshInstance: 306 def __init__(self): 307 self.mesh_ptr = _ti_core.create_mesh() 308 self.relation_set = set() 309 self.verts = MeshElementField(self, MeshElementType.Vertex, {}, {}, {}) 310 self.edges = MeshElementField(self, MeshElementType.Edge, {}, {}, {}) 311 self.faces = MeshElementField(self, MeshElementType.Face, {}, {}, {}) 312 self.cells = MeshElementField(self, MeshElementType.Cell, {}, {}, {}) 313 314 def get_position_as_numpy(self): 315 """Get the vertex position of current mesh to numpy array. 316 317 Returns: 318 3d numpy array: [x, y, z] with float-format. 319 """ 320 if hasattr(self, "_vert_position"): 321 return self._vert_position 322 raise TaichiSyntaxError("Position info is not in the file.") 323 324 def set_owned_offset(self, element_type: MeshElementType, owned_offset: ScalarField): 325 _ti_core.set_owned_offset(self.mesh_ptr, element_type, owned_offset.vars[0].ptr.snode()) 326 327 def set_total_offset(self, element_type: MeshElementType, total_offset: ScalarField): 328 _ti_core.set_total_offset(self.mesh_ptr, element_type, total_offset.vars[0].ptr.snode()) 329 330 def set_index_mapping(self, element_type: MeshElementType, conv_type: ConvType, mapping: ScalarField): 331 _ti_core.set_index_mapping(self.mesh_ptr, element_type, conv_type, mapping.vars[0].ptr.snode()) 332 333 def set_num_patches(self, num_patches: int): 334 _ti_core.set_num_patches(self.mesh_ptr, num_patches) 335 336 def set_patch_max_element_num(self, element_type: MeshElementType, max_element_num: int): 337 _ti_core.set_patch_max_element_num(self.mesh_ptr, element_type, max_element_num) 338 339 def set_relation_fixed(self, rel_type: MeshRelationType, value: ScalarField): 340 self.relation_set.add(rel_type) 341 _ti_core.set_relation_fixed(self.mesh_ptr, rel_type, value.vars[0].ptr.snode()) 342 343 def set_relation_dynamic( 344 self, 345 rel_type: MeshRelationType, 346 value: ScalarField, 347 patch_offset: ScalarField, 348 offset: ScalarField, 349 ): 350 self.relation_set.add(rel_type) 351 _ti_core.set_relation_dynamic( 352 self.mesh_ptr, 353 rel_type, 354 value.vars[0].ptr.snode(), 355 patch_offset.vars[0].ptr.snode(), 356 offset.vars[0].ptr.snode(), 357 ) 358 359 def add_mesh_attribute(self, element_type, snode, reorder_type): 360 _ti_core.add_mesh_attribute(self.mesh_ptr, element_type, snode, reorder_type) 361 362 def get_relation_size(self, from_index, to_element_type): 363 return _ti_core.get_relation_size( 364 self.mesh_ptr, 365 from_index.ptr, 366 to_element_type, 367 _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()), 368 ) 369 370 def get_relation_access(self, from_index, to_element_type, neighbor_idx_ptr): 371 return _ti_core.get_relation_access( 372 self.mesh_ptr, 373 from_index.ptr, 374 to_element_type, 375 neighbor_idx_ptr, 376 _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()), 377 )
314 def get_position_as_numpy(self): 315 """Get the vertex position of current mesh to numpy array. 316 317 Returns: 318 3d numpy array: [x, y, z] with float-format. 319 """ 320 if hasattr(self, "_vert_position"): 321 return self._vert_position 322 raise TaichiSyntaxError("Position info is not in the file.")
Get the vertex position of current mesh to numpy array.
Returns: 3d numpy array: [x, y, z] with float-format.
343 def set_relation_dynamic( 344 self, 345 rel_type: MeshRelationType, 346 value: ScalarField, 347 patch_offset: ScalarField, 348 offset: ScalarField, 349 ): 350 self.relation_set.add(rel_type) 351 _ti_core.set_relation_dynamic( 352 self.mesh_ptr, 353 rel_type, 354 value.vars[0].ptr.snode(), 355 patch_offset.vars[0].ptr.snode(), 356 offset.vars[0].ptr.snode(), 357 )
21class Ndarray: 22 """Taichi ndarray class. 23 24 Args: 25 dtype (DataType): Data type of each value. 26 shape (Tuple[int]): Shape of the Ndarray. 27 """ 28 29 def __init__(self): 30 self.host_accessor = None 31 self.shape = None 32 self.element_type = None 33 self.dtype = None 34 self.arr = None 35 self.layout = Layout.AOS 36 self.grad: "ScalarNdarray | VectorNdarray | MatrixNdarray | None" = None 37 38 def get_type(self): 39 return NdarrayTypeMetadata(self.element_type, self.shape, self.grad is not None) 40 41 @property 42 def element_shape(self): 43 """Gets ndarray element shape. 44 45 Returns: 46 Tuple[Int]: Ndarray element shape. 47 """ 48 raise NotImplementedError() 49 50 @python_scope 51 def __setitem__(self, key, value): 52 """Sets ndarray element in Python scope. 53 54 Args: 55 key (Union[List[int], int, None]): Coordinates of the ndarray element. 56 value (element type): Value to set. 57 """ 58 raise NotImplementedError() 59 60 @python_scope 61 def __getitem__(self, key): 62 """Gets ndarray element in Python scope. 63 64 Args: 65 key (Union[List[int], int, None]): Coordinates of the ndarray element. 66 67 Returns: 68 element type: Value retrieved. 69 """ 70 raise NotImplementedError() 71 72 @python_scope 73 def fill(self, val): 74 """Fills ndarray with a specific scalar value. 75 76 Args: 77 val (Union[int, float]): Value to fill. 78 """ 79 if impl.current_cfg().arch != _ti_core.Arch.cuda and impl.current_cfg().arch != _ti_core.Arch.x64: 80 self._fill_by_kernel(val) 81 elif _ti_core.is_tensor(self.element_type): 82 self._fill_by_kernel(val) 83 elif self.dtype == primitive_types.f32: 84 impl.get_runtime().prog.fill_float(self.arr, val) 85 elif self.dtype == primitive_types.i32: 86 impl.get_runtime().prog.fill_int(self.arr, val) 87 elif self.dtype == primitive_types.u32: 88 impl.get_runtime().prog.fill_uint(self.arr, val) 89 else: 90 self._fill_by_kernel(val) 91 92 @python_scope 93 def _ndarray_to_numpy(self): 94 """Converts ndarray to a numpy array. 95 96 Returns: 97 numpy.ndarray: The result numpy array. 98 """ 99 arr = np.zeros(shape=self.arr.total_shape(), dtype=to_numpy_type(self.dtype)) 100 from taichi._kernels import ndarray_to_ext_arr # pylint: disable=C0415 101 102 ndarray_to_ext_arr(self, arr) 103 impl.get_runtime().sync() 104 return arr 105 106 @python_scope 107 def _ndarray_matrix_to_numpy(self, as_vector): 108 """Converts matrix ndarray to a numpy array. 109 110 Returns: 111 numpy.ndarray: The result numpy array. 112 """ 113 arr = np.zeros(shape=self.arr.total_shape(), dtype=to_numpy_type(self.dtype)) 114 from taichi._kernels import ndarray_matrix_to_ext_arr # pylint: disable=C0415 115 116 layout_is_aos = 1 117 ndarray_matrix_to_ext_arr(self, arr, layout_is_aos, as_vector) 118 impl.get_runtime().sync() 119 return arr 120 121 @python_scope 122 def _ndarray_from_numpy(self, arr): 123 """Loads all values from a numpy array. 124 125 Args: 126 arr (numpy.ndarray): The source numpy array. 127 """ 128 if not isinstance(arr, np.ndarray): 129 raise TypeError(f"{np.ndarray} expected, but {type(arr)} provided") 130 if tuple(self.arr.total_shape()) != tuple(arr.shape): 131 raise ValueError(f"Mismatch shape: {tuple(self.arr.shape)} expected, but {tuple(arr.shape)} provided") 132 if not arr.flags.c_contiguous: 133 arr = np.ascontiguousarray(arr) 134 135 from taichi._kernels import ext_arr_to_ndarray # pylint: disable=C0415 136 137 ext_arr_to_ndarray(arr, self) 138 impl.get_runtime().sync() 139 140 @python_scope 141 def _ndarray_matrix_from_numpy(self, arr, as_vector): 142 """Loads all values from a numpy array. 143 144 Args: 145 arr (numpy.ndarray): The source numpy array. 146 """ 147 if not isinstance(arr, np.ndarray): 148 raise TypeError(f"{np.ndarray} expected, but {type(arr)} provided") 149 if tuple(self.arr.total_shape()) != tuple(arr.shape): 150 raise ValueError( 151 f"Mismatch shape: {tuple(self.arr.total_shape())} expected, but {tuple(arr.shape)} provided" 152 ) 153 if not arr.flags.c_contiguous: 154 arr = np.ascontiguousarray(arr) 155 156 from taichi._kernels import ext_arr_to_ndarray_matrix # pylint: disable=C0415 157 158 layout_is_aos = 1 159 ext_arr_to_ndarray_matrix(arr, self, layout_is_aos, as_vector) 160 impl.get_runtime().sync() 161 162 @python_scope 163 def _get_element_size(self): 164 """Returns the size of one element in bytes. 165 166 Returns: 167 Size in bytes. 168 """ 169 return self.arr.element_size() 170 171 @python_scope 172 def _get_nelement(self): 173 """Returns the total number of elements. 174 175 Returns: 176 Total number of elements. 177 """ 178 return self.arr.nelement() 179 180 @python_scope 181 def copy_from(self, other): 182 """Copies all elements from another ndarray. 183 184 The shape of the other ndarray needs to be the same as `self`. 185 186 Args: 187 other (Ndarray): The source ndarray. 188 """ 189 assert isinstance(other, Ndarray) 190 assert tuple(self.arr.shape) == tuple(other.arr.shape) 191 from taichi._kernels import ndarray_to_ndarray # pylint: disable=C0415 192 193 ndarray_to_ndarray(self, other) 194 impl.get_runtime().sync() 195 196 def _set_grad(self, grad: "ScalarNdarray | VectorNdarray | MatrixNdarray"): 197 """Sets the gradient ndarray. 198 199 Args: 200 grad (Ndarray): The gradient ndarray. 201 """ 202 self.grad = grad 203 204 def __deepcopy__(self, memo=None): 205 """Copies all elements to a new ndarray. 206 207 Returns: 208 Ndarray: The result ndarray. 209 """ 210 raise NotImplementedError() 211 212 def _fill_by_kernel(self, val): 213 """Fills ndarray with a specific scalar value using a ti.kernel. 214 215 Args: 216 val (Union[int, float]): Value to fill. 217 """ 218 raise NotImplementedError() 219 220 @python_scope 221 def _pad_key(self, key): 222 if key is None: 223 key = () 224 if not isinstance(key, (tuple, list)): 225 key = (key,) 226 if len(key) != len(self.arr.total_shape()): 227 raise TaichiIndexError(f"{len(self.arr.total_shape())}d ndarray indexed with {len(key)}d indices: {key}") 228 return key 229 230 @python_scope 231 def _initialize_host_accessor(self): 232 if self.host_accessor: 233 return 234 impl.get_runtime().materialize() 235 self.host_accessor = NdarrayHostAccessor(self.arr)
Taichi ndarray class.
Args: dtype (DataType): Data type of each value. shape (Tuple[int]): Shape of the Ndarray.
41 @property 42 def element_shape(self): 43 """Gets ndarray element shape. 44 45 Returns: 46 Tuple[Int]: Ndarray element shape. 47 """ 48 raise NotImplementedError()
Gets ndarray element shape.
Returns: Tuple[Int]: Ndarray element shape.
72 @python_scope 73 def fill(self, val): 74 """Fills ndarray with a specific scalar value. 75 76 Args: 77 val (Union[int, float]): Value to fill. 78 """ 79 if impl.current_cfg().arch != _ti_core.Arch.cuda and impl.current_cfg().arch != _ti_core.Arch.x64: 80 self._fill_by_kernel(val) 81 elif _ti_core.is_tensor(self.element_type): 82 self._fill_by_kernel(val) 83 elif self.dtype == primitive_types.f32: 84 impl.get_runtime().prog.fill_float(self.arr, val) 85 elif self.dtype == primitive_types.i32: 86 impl.get_runtime().prog.fill_int(self.arr, val) 87 elif self.dtype == primitive_types.u32: 88 impl.get_runtime().prog.fill_uint(self.arr, val) 89 else: 90 self._fill_by_kernel(val)
Fills ndarray with a specific scalar value.
Args: val (Union[int, float]): Value to fill.
180 @python_scope 181 def copy_from(self, other): 182 """Copies all elements from another ndarray. 183 184 The shape of the other ndarray needs to be the same as `self`. 185 186 Args: 187 other (Ndarray): The source ndarray. 188 """ 189 assert isinstance(other, Ndarray) 190 assert tuple(self.arr.shape) == tuple(other.arr.shape) 191 from taichi._kernels import ndarray_to_ndarray # pylint: disable=C0415 192 193 ndarray_to_ndarray(self, other) 194 impl.get_runtime().sync()
Copies all elements from another ndarray.
The shape of the other ndarray needs to be the same as self.
Args: other (Ndarray): The source ndarray.
15class SNode: 16 """A Python-side SNode wrapper. 17 18 For more information on Taichi's SNode system, please check out 19 these references: 20 21 * https://docs.taichi-lang.org/docs/sparse 22 * https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf 23 24 Arg: 25 ptr (pointer): The C++ side SNode pointer. 26 """ 27 28 def __init__(self, ptr: SNodeCxx) -> None: 29 self.ptr = ptr 30 31 def dense(self, axes: list[Axis], dimensions: list[int] | int) -> "SNode": 32 """Adds a dense SNode as a child component of `self`. 33 34 Args: 35 axes (List[Axis]): Axes to activate. 36 dimensions (Union[List[int], int]): Shape of each axis. 37 38 Returns: 39 The added :class:`~taichi.lang.SNode` instance. 40 """ 41 if isinstance(dimensions, numbers.Number): 42 dimensions = [dimensions] * len(axes) 43 return SNode(self.ptr.dense(axes, dimensions, _ti_core.DebugInfo(get_traceback()))) 44 45 def pointer(self, axes: list[Axis], dimensions: list[int] | int) -> "SNode": 46 """Adds a pointer SNode as a child component of `self`. 47 48 Args: 49 axes (List[Axis]): Axes to activate. 50 dimensions (Union[List[int], int]): Shape of each axis. 51 52 Returns: 53 The added :class:`~taichi.lang.SNode` instance. 54 """ 55 if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse): 56 raise TaichiRuntimeError("Pointer SNode is not supported on this backend.") 57 if isinstance(dimensions, numbers.Number): 58 dimensions = [dimensions] * len(axes) 59 return SNode(self.ptr.pointer(axes, dimensions, _ti_core.DebugInfo(get_traceback()))) 60 61 @staticmethod 62 def _hash(axes, dimensions): 63 # original code is #def hash(self,axes, dimensions) without #@staticmethod before fix pylint R0201 64 """Not supported.""" 65 raise RuntimeError("hash not yet supported") 66 # if isinstance(dimensions, int): 67 # dimensions = [dimensions] * len(axes) 68 # return SNode(self.ptr.hash(axes, dimensions)) 69 70 def dynamic(self, axis: list[Axis], dimension: int, chunk_size: int | None = None) -> "SNode": 71 """Adds a dynamic SNode as a child component of `self`. 72 73 Args: 74 axis (List[Axis]): Axis to activate, must be 1. 75 dimension (int): Shape of the axis. 76 chunk_size (int): Chunk size. 77 78 Returns: 79 The added :class:`~taichi.lang.SNode` instance. 80 """ 81 if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse): 82 raise TaichiRuntimeError("Dynamic SNode is not supported on this backend.") 83 assert len(axis) == 1 84 if chunk_size is None: 85 chunk_size = dimension 86 return SNode(self.ptr.dynamic(axis[0], dimension, chunk_size, _ti_core.DebugInfo(get_traceback()))) 87 88 def bitmasked(self, axes: list[Axis], dimensions: list[int] | int) -> "SNode": 89 """Adds a bitmasked SNode as a child component of `self`. 90 91 Args: 92 axes (List[Axis]): Axes to activate. 93 dimensions (Union[List[int], int]): Shape of each axis. 94 95 Returns: 96 The added :class:`~taichi.lang.SNode` instance. 97 """ 98 if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse): 99 raise TaichiRuntimeError("Bitmasked SNode is not supported on this backend.") 100 if isinstance(dimensions, numbers.Number): 101 dimensions = [dimensions] * len(axes) 102 return SNode(self.ptr.bitmasked(axes, dimensions, _ti_core.DebugInfo(get_traceback()))) 103 104 def quant_array(self, axes: list[Axis], dimensions: list[int] | int, max_num_bits: int) -> "SNode": 105 """Adds a quant_array SNode as a child component of `self`. 106 107 Args: 108 axes (List[Axis]): Axes to activate. 109 dimensions (Union[List[int], int]): Shape of each axis. 110 max_num_bits (int): Maximum number of bits it can hold. 111 112 Returns: 113 The added :class:`~taichi.lang.SNode` instance. 114 """ 115 if isinstance(dimensions, numbers.Number): 116 dimensions = [dimensions] * len(axes) 117 return SNode(self.ptr.quant_array(axes, dimensions, max_num_bits, _ti_core.DebugInfo(get_traceback()))) 118 119 def place(self, *args, offset: numbers.Number | tuple[numbers.Number] | None = None) -> "SNode": 120 """Places a list of Taichi fields under the `self` container. 121 122 Args: 123 *args (List[ti.field]): A list of Taichi fields to place. 124 offset (Union[Number, tuple[Number]]): Offset of the field domain. 125 126 Returns: 127 The `self` container. 128 """ 129 if offset is None: 130 offset = () 131 if isinstance(offset, numbers.Number): 132 offset = (offset,) 133 134 for arg in args: 135 if isinstance(arg, BitpackedFields): 136 bit_struct_type = arg.bit_struct_type_builder.build() 137 bit_struct_snode = self.ptr.bit_struct(bit_struct_type, _ti_core.DebugInfo(get_traceback())) 138 for field, id_in_bit_struct in arg.fields: 139 bit_struct_snode.place(field, offset, id_in_bit_struct) 140 elif isinstance(arg, Field): 141 for var in arg._get_field_members(): 142 self.ptr.place(var.ptr, offset, -1) 143 elif isinstance(arg, list): 144 for x in arg: 145 self.place(x, offset=offset) 146 else: 147 raise ValueError(f"{arg} cannot be placed") 148 return self 149 150 def lazy_grad(self): 151 """Automatically place the adjoint fields following the layout of their primal fields. 152 153 Users don't need to specify ``needs_grad`` when they define scalar/vector/matrix fields (primal fields) using autodiff. 154 When all the primal fields are defined, using ``taichi.root.lazy_grad()`` could automatically generate 155 their corresponding adjoint fields (gradient field). 156 157 To know more details about primal, adjoint fields and ``lazy_grad()``, 158 please see Page 4 and Page 13-14 of DiffTaichi Paper: https://arxiv.org/pdf/1910.00935.pdf 159 """ 160 self.ptr.lazy_grad() 161 162 def lazy_dual(self): 163 """Automatically place the dual fields following the layout of their primal fields.""" 164 self.ptr.lazy_dual() 165 166 def _allocate_adjoint_checkbit(self): 167 """Automatically place the adjoint flag fields following the layout of their primal fields for global data access rule checker""" 168 self.ptr.allocate_adjoint_checkbit() 169 170 def parent(self, n=1): 171 """Gets an ancestor of `self` in the SNode tree. 172 173 Args: 174 n (int): the number of levels going up from `self`. 175 176 Returns: 177 Union[None, _Root, SNode]: The n-th parent of `self`. 178 """ 179 p = self.ptr 180 while p and n > 0: 181 p = p.parent 182 n -= 1 183 if p is None: 184 return None 185 186 if p.type == _ti_core.SNodeType.root: 187 return impl.root 188 189 return SNode(p) 190 191 def _path_from_root(self): 192 """Gets the path from root to `self` in the SNode tree. 193 194 Returns: 195 List[Union[_Root, SNode]]: The list of SNodes on the path from root to `self`. 196 """ 197 p = self 198 res = [p] 199 while p != impl.root: 200 p = p.parent() 201 res.append(p) 202 res.reverse() 203 return res 204 205 @property 206 def _dtype(self): 207 """Gets the data type of `self`. 208 209 Returns: 210 DataType: The data type of `self`. 211 """ 212 return self.ptr.data_type() 213 214 @property 215 def _id(self): 216 """Gets the id of `self`. 217 218 Returns: 219 int: The id of `self`. 220 """ 221 return self.ptr.id 222 223 @property 224 def _snode_tree_id(self): 225 return self.ptr.get_snode_tree_id() 226 227 @property 228 def shape(self): 229 """Gets the number of elements from root in each axis of `self`. 230 231 Returns: 232 Tuple[int]: The number of elements from root in each axis of `self`. 233 """ 234 dim = self.ptr.num_active_indices() 235 ret = tuple(self.ptr.get_shape_along_axis(i) for i in range(dim)) 236 237 return ret 238 239 def _loop_range(self): 240 """Gets the taichi_python.SNode to serve as loop range. 241 242 Returns: 243 taichi_python.SNode: See above. 244 """ 245 return self.ptr 246 247 @property 248 def _name(self): 249 """Gets the name of `self`. 250 251 Returns: 252 str: The name of `self`. 253 """ 254 return self.ptr.name() 255 256 @property 257 def _snode(self): 258 """Gets `self`. 259 Returns: 260 SNode: `self`. 261 """ 262 return self 263 264 def _get_children(self): 265 """Gets all children components of `self`. 266 267 Returns: 268 List[SNode]: All children components of `self`. 269 """ 270 children = [] 271 for i in range(self.ptr.get_num_ch()): 272 children.append(SNode(self.ptr.get_ch(i))) 273 return children 274 275 @property 276 def _num_dynamically_allocated(self): 277 runtime = impl.get_runtime() 278 runtime.materialize_root_fb(False) 279 return runtime.prog.get_snode_num_dynamically_allocated(self.ptr) 280 281 @property 282 def _cell_size_bytes(self): 283 impl.get_runtime().materialize_root_fb(False) 284 return self.ptr.cell_size_bytes 285 286 @property 287 def _offset_bytes_in_parent_cell(self): 288 impl.get_runtime().materialize_root_fb(False) 289 return self.ptr.offset_bytes_in_parent_cell 290 291 def deactivate_all(self): 292 """Recursively deactivate all children components of `self`.""" 293 ch = self._get_children() 294 for c in ch: 295 c.deactivate_all() 296 SNodeType = _ti_core.SNodeType 297 if self.ptr.type == SNodeType.pointer or self.ptr.type == SNodeType.bitmasked: 298 from taichi._kernels import snode_deactivate # pylint: disable=C0415 299 300 snode_deactivate(self) 301 if self.ptr.type == SNodeType.dynamic: 302 # Note that dynamic nodes are different from other sparse nodes: 303 # instead of deactivating each element, we only need to deactivate 304 # its parent, whose linked list of chunks of elements will be deleted. 305 from taichi._kernels import ( # pylint: disable=C0415 306 snode_deactivate_dynamic, 307 ) 308 309 snode_deactivate_dynamic(self) 310 311 def __repr__(self): 312 type_ = str(self.ptr.type)[len("SNodeType.") :] 313 return f"<ti.SNode of type {type_}>" 314 315 def __str__(self): 316 # ti.root.dense(ti.i, 3).dense(ti.jk, (4, 5)).place(x) 317 # ti.root => dense [3] => dense [3, 4, 5] => place [3, 4, 5] 318 type_ = str(self.ptr.type)[len("SNodeType.") :] 319 shape = str(list(self.shape)) 320 parent = str(self.parent()) 321 return f"{parent} => {type_} {shape}" 322 323 def __eq__(self, other): 324 return self.ptr == other.ptr 325 326 def _physical_index_position(self): 327 """Gets mappings from virtual axes to physical axes. 328 329 Returns: 330 Dict[int, int]: Mappings from virtual axes to physical axes. 331 """ 332 ret = {} 333 for virtual, physical in enumerate(self.ptr.get_physical_index_position()): 334 if physical != -1: 335 ret[virtual] = physical 336 return ret
A Python-side SNode wrapper.
For more information on Taichi's SNode system, please check out these references:
- https://docs.taichi-lang.org/docs/sparse
- https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf
Arg: ptr (pointer): The C++ side SNode pointer.
31 def dense(self, axes: list[Axis], dimensions: list[int] | int) -> "SNode": 32 """Adds a dense SNode as a child component of `self`. 33 34 Args: 35 axes (List[Axis]): Axes to activate. 36 dimensions (Union[List[int], int]): Shape of each axis. 37 38 Returns: 39 The added :class:`~taichi.lang.SNode` instance. 40 """ 41 if isinstance(dimensions, numbers.Number): 42 dimensions = [dimensions] * len(axes) 43 return SNode(self.ptr.dense(axes, dimensions, _ti_core.DebugInfo(get_traceback())))
Adds a dense SNode as a child component of self.
Args: axes (List[Axis]): Axes to activate. dimensions (Union[List[int], int]): Shape of each axis.
Returns:
The added ~taichi.lang.SNode instance.
45 def pointer(self, axes: list[Axis], dimensions: list[int] | int) -> "SNode": 46 """Adds a pointer SNode as a child component of `self`. 47 48 Args: 49 axes (List[Axis]): Axes to activate. 50 dimensions (Union[List[int], int]): Shape of each axis. 51 52 Returns: 53 The added :class:`~taichi.lang.SNode` instance. 54 """ 55 if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse): 56 raise TaichiRuntimeError("Pointer SNode is not supported on this backend.") 57 if isinstance(dimensions, numbers.Number): 58 dimensions = [dimensions] * len(axes) 59 return SNode(self.ptr.pointer(axes, dimensions, _ti_core.DebugInfo(get_traceback())))
Adds a pointer SNode as a child component of self.
Args: axes (List[Axis]): Axes to activate. dimensions (Union[List[int], int]): Shape of each axis.
Returns:
The added ~taichi.lang.SNode instance.
70 def dynamic(self, axis: list[Axis], dimension: int, chunk_size: int | None = None) -> "SNode": 71 """Adds a dynamic SNode as a child component of `self`. 72 73 Args: 74 axis (List[Axis]): Axis to activate, must be 1. 75 dimension (int): Shape of the axis. 76 chunk_size (int): Chunk size. 77 78 Returns: 79 The added :class:`~taichi.lang.SNode` instance. 80 """ 81 if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse): 82 raise TaichiRuntimeError("Dynamic SNode is not supported on this backend.") 83 assert len(axis) == 1 84 if chunk_size is None: 85 chunk_size = dimension 86 return SNode(self.ptr.dynamic(axis[0], dimension, chunk_size, _ti_core.DebugInfo(get_traceback())))
Adds a dynamic SNode as a child component of self.
Args: axis (List[Axis]): Axis to activate, must be 1. dimension (int): Shape of the axis. chunk_size (int): Chunk size.
Returns:
The added ~taichi.lang.SNode instance.
88 def bitmasked(self, axes: list[Axis], dimensions: list[int] | int) -> "SNode": 89 """Adds a bitmasked SNode as a child component of `self`. 90 91 Args: 92 axes (List[Axis]): Axes to activate. 93 dimensions (Union[List[int], int]): Shape of each axis. 94 95 Returns: 96 The added :class:`~taichi.lang.SNode` instance. 97 """ 98 if not _ti_core.is_extension_supported(impl.current_cfg().arch, _ti_core.Extension.sparse): 99 raise TaichiRuntimeError("Bitmasked SNode is not supported on this backend.") 100 if isinstance(dimensions, numbers.Number): 101 dimensions = [dimensions] * len(axes) 102 return SNode(self.ptr.bitmasked(axes, dimensions, _ti_core.DebugInfo(get_traceback())))
Adds a bitmasked SNode as a child component of self.
Args: axes (List[Axis]): Axes to activate. dimensions (Union[List[int], int]): Shape of each axis.
Returns:
The added ~taichi.lang.SNode instance.
104 def quant_array(self, axes: list[Axis], dimensions: list[int] | int, max_num_bits: int) -> "SNode": 105 """Adds a quant_array SNode as a child component of `self`. 106 107 Args: 108 axes (List[Axis]): Axes to activate. 109 dimensions (Union[List[int], int]): Shape of each axis. 110 max_num_bits (int): Maximum number of bits it can hold. 111 112 Returns: 113 The added :class:`~taichi.lang.SNode` instance. 114 """ 115 if isinstance(dimensions, numbers.Number): 116 dimensions = [dimensions] * len(axes) 117 return SNode(self.ptr.quant_array(axes, dimensions, max_num_bits, _ti_core.DebugInfo(get_traceback())))
Adds a quant_array SNode as a child component of self.
Args: axes (List[Axis]): Axes to activate. dimensions (Union[List[int], int]): Shape of each axis. max_num_bits (int): Maximum number of bits it can hold.
Returns:
The added ~taichi.lang.SNode instance.
119 def place(self, *args, offset: numbers.Number | tuple[numbers.Number] | None = None) -> "SNode": 120 """Places a list of Taichi fields under the `self` container. 121 122 Args: 123 *args (List[ti.field]): A list of Taichi fields to place. 124 offset (Union[Number, tuple[Number]]): Offset of the field domain. 125 126 Returns: 127 The `self` container. 128 """ 129 if offset is None: 130 offset = () 131 if isinstance(offset, numbers.Number): 132 offset = (offset,) 133 134 for arg in args: 135 if isinstance(arg, BitpackedFields): 136 bit_struct_type = arg.bit_struct_type_builder.build() 137 bit_struct_snode = self.ptr.bit_struct(bit_struct_type, _ti_core.DebugInfo(get_traceback())) 138 for field, id_in_bit_struct in arg.fields: 139 bit_struct_snode.place(field, offset, id_in_bit_struct) 140 elif isinstance(arg, Field): 141 for var in arg._get_field_members(): 142 self.ptr.place(var.ptr, offset, -1) 143 elif isinstance(arg, list): 144 for x in arg: 145 self.place(x, offset=offset) 146 else: 147 raise ValueError(f"{arg} cannot be placed") 148 return self
Places a list of Taichi fields under the self container.
Args: *args (List[ti.field]): A list of Taichi fields to place. offset (Union[Number, tuple[Number]]): Offset of the field domain.
Returns:
The self container.
150 def lazy_grad(self): 151 """Automatically place the adjoint fields following the layout of their primal fields. 152 153 Users don't need to specify ``needs_grad`` when they define scalar/vector/matrix fields (primal fields) using autodiff. 154 When all the primal fields are defined, using ``taichi.root.lazy_grad()`` could automatically generate 155 their corresponding adjoint fields (gradient field). 156 157 To know more details about primal, adjoint fields and ``lazy_grad()``, 158 please see Page 4 and Page 13-14 of DiffTaichi Paper: https://arxiv.org/pdf/1910.00935.pdf 159 """ 160 self.ptr.lazy_grad()
Automatically place the adjoint fields following the layout of their primal fields.
Users don't need to specify needs_grad when they define scalar/vector/matrix fields (primal fields) using autodiff.
When all the primal fields are defined, using taichi.root.lazy_grad() could automatically generate
their corresponding adjoint fields (gradient field).
To know more details about primal, adjoint fields and lazy_grad(),
please see Page 4 and Page 13-14 of DiffTaichi Paper: https://arxiv.org/pdf/1910.00935.pdf
162 def lazy_dual(self): 163 """Automatically place the dual fields following the layout of their primal fields.""" 164 self.ptr.lazy_dual()
Automatically place the dual fields following the layout of their primal fields.
170 def parent(self, n=1): 171 """Gets an ancestor of `self` in the SNode tree. 172 173 Args: 174 n (int): the number of levels going up from `self`. 175 176 Returns: 177 Union[None, _Root, SNode]: The n-th parent of `self`. 178 """ 179 p = self.ptr 180 while p and n > 0: 181 p = p.parent 182 n -= 1 183 if p is None: 184 return None 185 186 if p.type == _ti_core.SNodeType.root: 187 return impl.root 188 189 return SNode(p)
Gets an ancestor of self in the SNode tree.
Args:
n (int): the number of levels going up from self.
Returns:
Union[None, _Root, SNode]: The n-th parent of self.
227 @property 228 def shape(self): 229 """Gets the number of elements from root in each axis of `self`. 230 231 Returns: 232 Tuple[int]: The number of elements from root in each axis of `self`. 233 """ 234 dim = self.ptr.num_active_indices() 235 ret = tuple(self.ptr.get_shape_along_axis(i) for i in range(dim)) 236 237 return ret
Gets the number of elements from root in each axis of self.
Returns:
Tuple[int]: The number of elements from root in each axis of self.
291 def deactivate_all(self): 292 """Recursively deactivate all children components of `self`.""" 293 ch = self._get_children() 294 for c in ch: 295 c.deactivate_all() 296 SNodeType = _ti_core.SNodeType 297 if self.ptr.type == SNodeType.pointer or self.ptr.type == SNodeType.bitmasked: 298 from taichi._kernels import snode_deactivate # pylint: disable=C0415 299 300 snode_deactivate(self) 301 if self.ptr.type == SNodeType.dynamic: 302 # Note that dynamic nodes are different from other sparse nodes: 303 # instead of deactivating each element, we only need to deactivate 304 # its parent, whose linked list of chunks of elements will be deleted. 305 from taichi._kernels import ( # pylint: disable=C0415 306 snode_deactivate_dynamic, 307 ) 308 309 snode_deactivate_dynamic(self)
Recursively deactivate all children components of self.
275class ScalarField(Field): 276 """Taichi scalar field with SNode implementation. 277 278 Args: 279 var (Expr): Field member. 280 """ 281 282 def __init__(self, var): 283 super().__init__([var]) 284 285 def fill(self, val): 286 """Fills this scalar field with a specified value.""" 287 if in_python_scope(): 288 from taichi._kernels import fill_field # pylint: disable=C0415 289 290 fill_field(self, val) 291 else: 292 from taichi._funcs import field_fill_taichi_scope # pylint: disable=C0415 293 294 field_fill_taichi_scope(self, val) 295 296 @python_scope 297 def to_numpy(self, dtype=None): 298 """Converts this field to a `numpy.ndarray`.""" 299 if self.parent()._snode.ptr.type == _ti_core.SNodeType.dynamic: 300 warn( 301 "You are trying to convert a dynamic snode to a numpy array, be aware that inactive items in the snode will be converted to zeros in the resulting array." 302 ) 303 if dtype is None: 304 dtype = to_numpy_type(self.dtype) 305 import numpy as np # pylint: disable=C0415 306 307 arr = np.zeros(shape=self.shape, dtype=dtype) 308 from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415 309 310 tensor_to_ext_arr(self, arr) 311 taichi.lang.runtime_ops.sync() 312 return arr 313 314 @python_scope 315 def to_torch(self, device=None): 316 """Converts this field to a `torch.tensor`.""" 317 import torch # pylint: disable=C0415 318 319 # pylint: disable=E1101 320 arr = torch.zeros(size=self.shape, dtype=to_pytorch_type(self.dtype), device=device) 321 from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415 322 323 tensor_to_ext_arr(self, arr) 324 taichi.lang.runtime_ops.sync() 325 return arr 326 327 @python_scope 328 def to_paddle(self, place=None): 329 """Converts this field to a `paddle.Tensor`.""" 330 import paddle # pylint: disable=C0415 331 332 # pylint: disable=E1101 333 # paddle.empty() doesn't support argument `place`` 334 arr = paddle.to_tensor(paddle.zeros(self.shape, to_paddle_type(self.dtype)), place=place) 335 from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415 336 337 tensor_to_ext_arr(self, arr) 338 taichi.lang.runtime_ops.sync() 339 return arr 340 341 @python_scope 342 def _from_external_arr(self, arr): 343 if len(self.shape) != len(arr.shape): 344 raise ValueError(f"ti.field shape {self.shape} does not match" f" the numpy array shape {arr.shape}") 345 for i, _ in enumerate(self.shape): 346 if self.shape[i] != arr.shape[i]: 347 raise ValueError(f"ti.field shape {self.shape} does not match" f" the numpy array shape {arr.shape}") 348 from taichi._kernels import ext_arr_to_tensor # pylint: disable=C0415 349 350 ext_arr_to_tensor(arr, self) 351 taichi.lang.runtime_ops.sync() 352 353 @python_scope 354 def from_numpy(self, arr): 355 """Copies the data from a `numpy.ndarray` into this field.""" 356 if not arr.flags.c_contiguous: 357 import numpy as np # pylint: disable=C0415 358 359 arr = np.ascontiguousarray(arr) 360 self._from_external_arr(arr) 361 362 @python_scope 363 def __setitem__(self, key, value): 364 self._initialize_host_accessors() 365 self.host_accessors[0].setter(value, *self._pad_key(key)) 366 367 @python_scope 368 def __getitem__(self, key): 369 self._initialize_host_accessors() 370 # Check for potential slicing behaviour 371 # for instance: x[0, :] 372 padded_key = self._pad_key(key) 373 import numpy as np # pylint: disable=C0415 374 375 for key in padded_key: 376 if not isinstance(key, (int, np.integer)): 377 raise TypeError( 378 f"Detected illegal element of type: {type(key)}. " 379 f"Please be aware that slicing a ti.field is not supported so far." 380 ) 381 return self.host_accessors[0].getter(*padded_key) 382 383 def __repr__(self): 384 # make interactive shell happy, prevent materialization 385 return "<ti.field>"
Taichi scalar field with SNode implementation.
Args: var (Expr): Field member.
285 def fill(self, val): 286 """Fills this scalar field with a specified value.""" 287 if in_python_scope(): 288 from taichi._kernels import fill_field # pylint: disable=C0415 289 290 fill_field(self, val) 291 else: 292 from taichi._funcs import field_fill_taichi_scope # pylint: disable=C0415 293 294 field_fill_taichi_scope(self, val)
Fills this scalar field with a specified value.
296 @python_scope 297 def to_numpy(self, dtype=None): 298 """Converts this field to a `numpy.ndarray`.""" 299 if self.parent()._snode.ptr.type == _ti_core.SNodeType.dynamic: 300 warn( 301 "You are trying to convert a dynamic snode to a numpy array, be aware that inactive items in the snode will be converted to zeros in the resulting array." 302 ) 303 if dtype is None: 304 dtype = to_numpy_type(self.dtype) 305 import numpy as np # pylint: disable=C0415 306 307 arr = np.zeros(shape=self.shape, dtype=dtype) 308 from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415 309 310 tensor_to_ext_arr(self, arr) 311 taichi.lang.runtime_ops.sync() 312 return arr
Converts this field to a numpy.ndarray.
314 @python_scope 315 def to_torch(self, device=None): 316 """Converts this field to a `torch.tensor`.""" 317 import torch # pylint: disable=C0415 318 319 # pylint: disable=E1101 320 arr = torch.zeros(size=self.shape, dtype=to_pytorch_type(self.dtype), device=device) 321 from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415 322 323 tensor_to_ext_arr(self, arr) 324 taichi.lang.runtime_ops.sync() 325 return arr
Converts this field to a torch.tensor.
327 @python_scope 328 def to_paddle(self, place=None): 329 """Converts this field to a `paddle.Tensor`.""" 330 import paddle # pylint: disable=C0415 331 332 # pylint: disable=E1101 333 # paddle.empty() doesn't support argument `place`` 334 arr = paddle.to_tensor(paddle.zeros(self.shape, to_paddle_type(self.dtype)), place=place) 335 from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415 336 337 tensor_to_ext_arr(self, arr) 338 taichi.lang.runtime_ops.sync() 339 return arr
Converts this field to a paddle.Tensor.
353 @python_scope 354 def from_numpy(self, arr): 355 """Copies the data from a `numpy.ndarray` into this field.""" 356 if not arr.flags.c_contiguous: 357 import numpy as np # pylint: disable=C0415 358 359 arr = np.ascontiguousarray(arr) 360 self._from_external_arr(arr)
Copies the data from a numpy.ndarray into this field.
Inherited Members
238class ScalarNdarray(Ndarray): 239 """Taichi ndarray with scalar elements. 240 241 Args: 242 dtype (DataType): Data type of each value. 243 shape (Tuple[int]): Shape of the ndarray. 244 """ 245 246 def __init__(self, dtype, arr_shape): 247 super().__init__() 248 self.dtype = cook_dtype(dtype) 249 self.arr = impl.get_runtime().prog.create_ndarray( 250 self.dtype, arr_shape, layout=Layout.NULL, zero_fill=True, dbg_info=_ti_core.DebugInfo(get_traceback()) 251 ) 252 self.shape = tuple(self.arr.shape) 253 self.element_type = dtype 254 255 def __del__(self): 256 if ( 257 impl is not None 258 and impl.get_runtime is not None 259 and impl.get_runtime() is not None 260 and impl.get_runtime().prog is not None 261 ): 262 impl.get_runtime().prog.delete_ndarray(self.arr) 263 264 @property 265 def element_shape(self): 266 return () 267 268 @python_scope 269 def __setitem__(self, key, value): 270 self._initialize_host_accessor() 271 self.host_accessor.setter(value, *self._pad_key(key)) 272 273 @python_scope 274 def __getitem__(self, key): 275 self._initialize_host_accessor() 276 return self.host_accessor.getter(*self._pad_key(key)) 277 278 @python_scope 279 def to_numpy(self): 280 return self._ndarray_to_numpy() 281 282 @python_scope 283 def from_numpy(self, arr): 284 self._ndarray_from_numpy(arr) 285 286 def __deepcopy__(self, memo=None): 287 ret_arr = ScalarNdarray(self.dtype, self.shape) 288 ret_arr.copy_from(self) 289 return ret_arr 290 291 def _fill_by_kernel(self, val): 292 from taichi._kernels import fill_ndarray # pylint: disable=C0415 293 294 fill_ndarray(self, val) 295 296 def __repr__(self): 297 return "<ti.ndarray>"
Taichi ndarray with scalar elements.
Args: dtype (DataType): Data type of each value. shape (Tuple[int]): Shape of the ndarray.
246 def __init__(self, dtype, arr_shape): 247 super().__init__() 248 self.dtype = cook_dtype(dtype) 249 self.arr = impl.get_runtime().prog.create_ndarray( 250 self.dtype, arr_shape, layout=Layout.NULL, zero_fill=True, dbg_info=_ti_core.DebugInfo(get_traceback()) 251 ) 252 self.shape = tuple(self.arr.shape) 253 self.element_type = dtype
26class Struct: 27 """The Struct type class. 28 29 A struct is a dictionary-like data structure that stores members as 30 (key, value) pairs. Valid data members of a struct can be scalars, 31 matrices or other dictionary-like structures. 32 33 Args: 34 entries (Dict[str, Union[Dict, Expr, Matrix, Struct]]): \ 35 keys and values for struct members. Entries can optionally 36 include a dictionary of functions with the key '__struct_methods' 37 which will be attached to the struct for executing on the struct data. 38 39 Returns: 40 An instance of this struct. 41 42 Example:: 43_ 44 >>> vec3 = ti.types.vector(3, ti.f32) 45 >>> a = ti.Struct(v=vec3([0, 0, 0]), t=1.0) 46 >>> print(a.items) 47 dict_items([('v', [0. 0. 0.]), ('t', 1.0)]) 48 >>> 49 >>> B = ti.Struct(v=vec3([0., 0., 0.]), t=1.0, A=a) 50 >>> print(B.items) 51 dict_items([('v', [0. 0. 0.]), ('t', 1.0), ('A', {'v': [[0.], [0.], [0.]], 't': 1.0})]) 52 """ 53 54 _is_taichi_class = True 55 _instance_count = 0 56 57 def __init__(self, *args, **kwargs): 58 # converts lists to matrices and dicts to structs 59 if len(args) == 1 and kwargs == {} and isinstance(args[0], dict): 60 self.__entries = args[0] 61 elif len(args) == 0: 62 self.__entries = kwargs 63 else: 64 raise TaichiSyntaxError( 65 "Custom structs need to be initialized using either dictionary or keyword arguments" 66 ) 67 self.__methods = self.__entries.pop("__struct_methods", {}) 68 matrix_ndim = self.__entries.pop("__matrix_ndim", {}) 69 self._register_methods() 70 71 for k, v in self.__entries.items(): 72 if isinstance(v, (list, tuple)): 73 v = Matrix(v) 74 if isinstance(v, dict): 75 v = Struct(v) 76 self.__entries[k] = v if in_python_scope() else impl.expr_init(v) 77 self._register_members() 78 self.__dtype = None 79 80 @property 81 def keys(self): 82 """Returns the list of member names in string format. 83 84 Example:: 85 86 >>> vec3 = ti.types.vector(3, ti.f32) 87 >>> sphere = ti.Struct(center=vec3([0, 0, 0]), radius=1.0) 88 >>> a.keys 89 ['center', 'radius'] 90 """ 91 return list(self.__entries.keys()) 92 93 @property 94 def _members(self): 95 return list(self.__entries.values()) 96 97 @property 98 def entries(self): 99 return self.__entries 100 101 @property 102 def methods(self): 103 return self.__methods 104 105 @property 106 def items(self): 107 """Returns the items in this struct. 108 109 Example:: 110 111 >>> vec3 = ti.types.vector(3, ti.f32) 112 >>> sphere = ti.Struct(center=vec3([0, 0, 0]), radius=1.0) 113 >>> sphere.items 114 dict_items([('center', 2), ('radius', 1.0)]) 115 """ 116 return self.__entries.items() 117 118 def _register_members(self): 119 # https://stackoverflow.com/questions/48448074/adding-a-property-to-an-existing-object-instance 120 cls = self.__class__ 121 new_cls_name = cls.__name__ + str(cls._instance_count) 122 cls._instance_count += 1 123 properties = {k: property(cls._make_getter(k), cls._make_setter(k)) for k in self.keys} 124 self.__class__ = type(new_cls_name, (cls,), properties) 125 126 def _register_methods(self): 127 for name, method in self.__methods.items(): 128 # use MethodType to pass self (this object) to the method 129 setattr(self, name, MethodType(method, self)) 130 131 def __getitem__(self, key): 132 ret = self.__entries[key] 133 if isinstance(ret, SNodeHostAccess): 134 ret = ret.accessor.getter(*ret.key) 135 return ret 136 137 def __setitem__(self, key, value): 138 if isinstance(self.__entries[key], SNodeHostAccess): 139 self.__entries[key].accessor.setter(value, *self.__entries[key].key) 140 else: 141 if in_python_scope(): 142 if isinstance(self.__entries[key], Struct) or isinstance(self.__entries[key], Matrix): 143 self.__entries[key]._set_entries(value) 144 else: 145 if isinstance(value, numbers.Number): 146 self.__entries[key] = value 147 else: 148 raise TypeError("A number is expected when assigning struct members") 149 else: 150 self.__entries[key] = value 151 152 def _set_entries(self, value): 153 if isinstance(value, dict): 154 value = Struct(value) 155 for k in self.keys: 156 self[k] = value[k] 157 self.__dtype = value.__dtype 158 159 @staticmethod 160 def _make_getter(key): 161 def getter(self): 162 """Get an entry from custom struct by name.""" 163 return self[key] 164 165 return getter 166 167 @staticmethod 168 def _make_setter(key): 169 @python_scope 170 def setter(self, value): 171 self[key] = value 172 173 return setter 174 175 @taichi_scope 176 def _assign(self, other): 177 if not isinstance(other, (dict, Struct)): 178 raise TaichiTypeError("Only dict or Struct can be assigned to a Struct") 179 if isinstance(other, dict): 180 other = Struct(other) 181 if self.__entries.keys() != other.__entries.keys(): 182 raise TaichiTypeError(f"Member mismatch between structs {self.keys}, {other.keys}") 183 for k, v in self.items: 184 v._assign(other.__entries[k]) 185 self.__dtype = other.__dtype 186 return self 187 188 def __len__(self): 189 """Get the number of entries in a custom struct""" 190 return len(self.__entries) 191 192 def __iter__(self): 193 return self.__entries.values() 194 195 def __str__(self): 196 """Python scope struct array print support.""" 197 if impl.inside_kernel(): 198 item_str = ", ".join([str(k) + "=" + str(v) for k, v in self.items]) 199 item_str += f", struct_methods={self.__methods}" 200 return f"<ti.Struct {item_str}>" 201 return str(self.to_dict()) 202 203 def __repr__(self): 204 return str(self.to_dict()) 205 206 def to_dict(self, include_methods=False, include_ndim=False): 207 """Converts the Struct to a dictionary. 208 209 Args: 210 include_methods (bool): Whether any struct methods should be included 211 in the result dictionary under the key '__struct_methods'. 212 213 Returns: 214 Dict: The result dictionary. 215 """ 216 res_dict = { 217 k: ( 218 v.to_dict(include_methods=include_methods, include_ndim=include_ndim) 219 if isinstance(v, Struct) 220 else v.to_list() if isinstance(v, Matrix) else v 221 ) 222 for k, v in self.__entries.items() 223 } 224 if include_methods: 225 res_dict["__struct_methods"] = self.__methods 226 if include_ndim: 227 res_dict["__matrix_ndim"] = dict() 228 for k, v in self.__entries.items(): 229 if isinstance(v, Matrix): 230 res_dict["__matrix_ndim"][k] = v.ndim 231 return res_dict 232 233 @classmethod 234 @python_scope 235 def field( 236 cls, 237 members, 238 methods={}, 239 shape=None, 240 name="<Struct>", 241 offset=None, 242 needs_grad=False, 243 needs_dual=False, 244 layout=Layout.AOS, 245 ): 246 """Creates a :class:`~taichi.StructField` with each element 247 has this struct as its type. 248 249 Args: 250 members (dict): a dict, each item is like `name: type`. 251 methods (dict): a dict of methods that should be included with 252 the field. Each struct item of the field will have the 253 methods as instance functions. 254 shape (Tuple[int]): width and height of the field. 255 offset (Tuple[int]): offset of the indices of the created field. 256 For example if `offset=(-10, -10)` the indices of the field 257 will start at `(-10, -10)`, not `(0, 0)`. 258 needs_grad (bool): enabling grad field (reverse mode autodiff) or not. 259 needs_dual (bool): enabling dual field (forward mode autodiff) or not. 260 layout: AOS or SOA. 261 262 Example: 263 264 >>> vec3 = ti.types.vector(3, ti.f32) 265 >>> sphere = {"center": vec3, "radius": float} 266 >>> F = ti.Struct.field(sphere, shape=(3, 3)) 267 >>> F 268 {'center': array([[[0., 0., 0.], 269 [0., 0., 0.], 270 [0., 0., 0.]], 271 272 [[0., 0., 0.], 273 [0., 0., 0.], 274 [0., 0., 0.]], 275 276 [[0., 0., 0.], 277 [0., 0., 0.], 278 [0., 0., 0.]]], dtype=float32), 'radius': array([[0., 0., 0.], 279 [0., 0., 0.], 280 [0., 0., 0.]], dtype=float32)} 281 """ 282 283 if shape is None and offset is not None: 284 raise TaichiSyntaxError("shape cannot be None when offset is being set") 285 286 field_dict = {} 287 288 for key, dtype in members.items(): 289 field_name = name + "." + key 290 if isinstance(dtype, CompoundType): 291 if isinstance(dtype, StructType): 292 field_dict[key] = dtype.field( 293 shape=None, 294 name=field_name, 295 offset=offset, 296 needs_grad=needs_grad, 297 needs_dual=needs_dual, 298 ) 299 else: 300 field_dict[key] = dtype.field( 301 shape=None, 302 name=field_name, 303 offset=offset, 304 needs_grad=needs_grad, 305 needs_dual=needs_dual, 306 ndim=getattr(dtype, "ndim", 2), 307 ) 308 else: 309 field_dict[key] = impl.field( 310 dtype, 311 shape=None, 312 name=field_name, 313 offset=offset, 314 needs_grad=needs_grad, 315 needs_dual=needs_dual, 316 ) 317 318 if shape is not None: 319 if isinstance(shape, numbers.Number): 320 shape = (shape,) 321 if isinstance(offset, numbers.Number): 322 offset = (offset,) 323 324 if offset is not None and len(shape) != len(offset): 325 raise TaichiSyntaxError( 326 f"The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})" 327 ) 328 dim = len(shape) 329 if layout == Layout.SOA: 330 for e in field_dict.values(): 331 impl.root.dense(impl.index_nd(dim), shape).place(e, offset=offset) 332 if needs_grad: 333 for e in field_dict.values(): 334 impl.root.dense(impl.index_nd(dim), shape).place(e.grad, offset=offset) 335 if needs_dual: 336 for e in field_dict.values(): 337 impl.root.dense(impl.index_nd(dim), shape).place(e.dual, offset=offset) 338 else: 339 impl.root.dense(impl.index_nd(dim), shape).place(*tuple(field_dict.values()), offset=offset) 340 if needs_grad: 341 grads = tuple(e.grad for e in field_dict.values()) 342 impl.root.dense(impl.index_nd(dim), shape).place(*grads, offset=offset) 343 344 if needs_dual: 345 duals = tuple(e.dual for e in field_dict.values()) 346 impl.root.dense(impl.index_nd(dim), shape).place(*duals, offset=offset) 347 348 return StructField(field_dict, methods, name=name)
The Struct type class.
A struct is a dictionary-like data structure that stores members as
(key, value) pairs. Valid data members of a struct can be scalars,
matrices or other dictionary-like structures.
Args:
entries (Dict[str, Union[Dict, Expr, Matrix, Struct]]): keys and values for struct members. Entries can optionally
include a dictionary of functions with the key '__struct_methods'
which will be attached to the struct for executing on the struct data.
Returns:
An instance of this struct.
Example::
_
vec3 = ti.types.vector(3, ti.f32) a = ti.Struct(v=vec3([0, 0, 0]), t=1.0) print(a.items) dict_items([('v', [0. 0. 0.]), ('t', 1.0)])
B = ti.Struct(v=vec3([0., 0., 0.]), t=1.0, A=a) print(B.items) dict_items([('v', [0. 0. 0.]), ('t', 1.0), ('A', {'v': [[0.], [0.], [0.]], 't': 1.0})])
57 def __init__(self, *args, **kwargs): 58 # converts lists to matrices and dicts to structs 59 if len(args) == 1 and kwargs == {} and isinstance(args[0], dict): 60 self.__entries = args[0] 61 elif len(args) == 0: 62 self.__entries = kwargs 63 else: 64 raise TaichiSyntaxError( 65 "Custom structs need to be initialized using either dictionary or keyword arguments" 66 ) 67 self.__methods = self.__entries.pop("__struct_methods", {}) 68 matrix_ndim = self.__entries.pop("__matrix_ndim", {}) 69 self._register_methods() 70 71 for k, v in self.__entries.items(): 72 if isinstance(v, (list, tuple)): 73 v = Matrix(v) 74 if isinstance(v, dict): 75 v = Struct(v) 76 self.__entries[k] = v if in_python_scope() else impl.expr_init(v) 77 self._register_members() 78 self.__dtype = None
80 @property 81 def keys(self): 82 """Returns the list of member names in string format. 83 84 Example:: 85 86 >>> vec3 = ti.types.vector(3, ti.f32) 87 >>> sphere = ti.Struct(center=vec3([0, 0, 0]), radius=1.0) 88 >>> a.keys 89 ['center', 'radius'] 90 """ 91 return list(self.__entries.keys())
Returns the list of member names in string format.
Example::
>>> vec3 = ti.types.vector(3, ti.f32)
>>> sphere = ti.Struct(center=vec3([0, 0, 0]), radius=1.0)
>>> a.keys
['center', 'radius']
105 @property 106 def items(self): 107 """Returns the items in this struct. 108 109 Example:: 110 111 >>> vec3 = ti.types.vector(3, ti.f32) 112 >>> sphere = ti.Struct(center=vec3([0, 0, 0]), radius=1.0) 113 >>> sphere.items 114 dict_items([('center', 2), ('radius', 1.0)]) 115 """ 116 return self.__entries.items()
Returns the items in this struct.
Example::
>>> vec3 = ti.types.vector(3, ti.f32)
>>> sphere = ti.Struct(center=vec3([0, 0, 0]), radius=1.0)
>>> sphere.items
dict_items([('center', 2), ('radius', 1.0)])
206 def to_dict(self, include_methods=False, include_ndim=False): 207 """Converts the Struct to a dictionary. 208 209 Args: 210 include_methods (bool): Whether any struct methods should be included 211 in the result dictionary under the key '__struct_methods'. 212 213 Returns: 214 Dict: The result dictionary. 215 """ 216 res_dict = { 217 k: ( 218 v.to_dict(include_methods=include_methods, include_ndim=include_ndim) 219 if isinstance(v, Struct) 220 else v.to_list() if isinstance(v, Matrix) else v 221 ) 222 for k, v in self.__entries.items() 223 } 224 if include_methods: 225 res_dict["__struct_methods"] = self.__methods 226 if include_ndim: 227 res_dict["__matrix_ndim"] = dict() 228 for k, v in self.__entries.items(): 229 if isinstance(v, Matrix): 230 res_dict["__matrix_ndim"][k] = v.ndim 231 return res_dict
Converts the Struct to a dictionary.
Args: include_methods (bool): Whether any struct methods should be included in the result dictionary under the key '__struct_methods'.
Returns: Dict: The result dictionary.
233 @classmethod 234 @python_scope 235 def field( 236 cls, 237 members, 238 methods={}, 239 shape=None, 240 name="<Struct>", 241 offset=None, 242 needs_grad=False, 243 needs_dual=False, 244 layout=Layout.AOS, 245 ): 246 """Creates a :class:`~taichi.StructField` with each element 247 has this struct as its type. 248 249 Args: 250 members (dict): a dict, each item is like `name: type`. 251 methods (dict): a dict of methods that should be included with 252 the field. Each struct item of the field will have the 253 methods as instance functions. 254 shape (Tuple[int]): width and height of the field. 255 offset (Tuple[int]): offset of the indices of the created field. 256 For example if `offset=(-10, -10)` the indices of the field 257 will start at `(-10, -10)`, not `(0, 0)`. 258 needs_grad (bool): enabling grad field (reverse mode autodiff) or not. 259 needs_dual (bool): enabling dual field (forward mode autodiff) or not. 260 layout: AOS or SOA. 261 262 Example: 263 264 >>> vec3 = ti.types.vector(3, ti.f32) 265 >>> sphere = {"center": vec3, "radius": float} 266 >>> F = ti.Struct.field(sphere, shape=(3, 3)) 267 >>> F 268 {'center': array([[[0., 0., 0.], 269 [0., 0., 0.], 270 [0., 0., 0.]], 271 272 [[0., 0., 0.], 273 [0., 0., 0.], 274 [0., 0., 0.]], 275 276 [[0., 0., 0.], 277 [0., 0., 0.], 278 [0., 0., 0.]]], dtype=float32), 'radius': array([[0., 0., 0.], 279 [0., 0., 0.], 280 [0., 0., 0.]], dtype=float32)} 281 """ 282 283 if shape is None and offset is not None: 284 raise TaichiSyntaxError("shape cannot be None when offset is being set") 285 286 field_dict = {} 287 288 for key, dtype in members.items(): 289 field_name = name + "." + key 290 if isinstance(dtype, CompoundType): 291 if isinstance(dtype, StructType): 292 field_dict[key] = dtype.field( 293 shape=None, 294 name=field_name, 295 offset=offset, 296 needs_grad=needs_grad, 297 needs_dual=needs_dual, 298 ) 299 else: 300 field_dict[key] = dtype.field( 301 shape=None, 302 name=field_name, 303 offset=offset, 304 needs_grad=needs_grad, 305 needs_dual=needs_dual, 306 ndim=getattr(dtype, "ndim", 2), 307 ) 308 else: 309 field_dict[key] = impl.field( 310 dtype, 311 shape=None, 312 name=field_name, 313 offset=offset, 314 needs_grad=needs_grad, 315 needs_dual=needs_dual, 316 ) 317 318 if shape is not None: 319 if isinstance(shape, numbers.Number): 320 shape = (shape,) 321 if isinstance(offset, numbers.Number): 322 offset = (offset,) 323 324 if offset is not None and len(shape) != len(offset): 325 raise TaichiSyntaxError( 326 f"The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})" 327 ) 328 dim = len(shape) 329 if layout == Layout.SOA: 330 for e in field_dict.values(): 331 impl.root.dense(impl.index_nd(dim), shape).place(e, offset=offset) 332 if needs_grad: 333 for e in field_dict.values(): 334 impl.root.dense(impl.index_nd(dim), shape).place(e.grad, offset=offset) 335 if needs_dual: 336 for e in field_dict.values(): 337 impl.root.dense(impl.index_nd(dim), shape).place(e.dual, offset=offset) 338 else: 339 impl.root.dense(impl.index_nd(dim), shape).place(*tuple(field_dict.values()), offset=offset) 340 if needs_grad: 341 grads = tuple(e.grad for e in field_dict.values()) 342 impl.root.dense(impl.index_nd(dim), shape).place(*grads, offset=offset) 343 344 if needs_dual: 345 duals = tuple(e.dual for e in field_dict.values()) 346 impl.root.dense(impl.index_nd(dim), shape).place(*duals, offset=offset) 347 348 return StructField(field_dict, methods, name=name)
Creates a ~taichi.StructField with each element
has this struct as its type.
Args:
members (dict): a dict, each item is like name: type.
methods (dict): a dict of methods that should be included with
the field. Each struct item of the field will have the
methods as instance functions.
shape (Tuple[int]): width and height of the field.
offset (Tuple[int]): offset of the indices of the created field.
For example if offset=(-10, -10) the indices of the field
will start at (-10, -10), not (0, 0).
needs_grad (bool): enabling grad field (reverse mode autodiff) or not.
needs_dual (bool): enabling dual field (forward mode autodiff) or not.
layout: AOS or SOA.
Example:
>>> vec3 = ti.types.vector(3, ti.f32)
>>> sphere = {"center": vec3, "radius": float}
>>> F = ti.Struct.field(sphere, shape=(3, 3))
>>> F
{'center': array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float32), 'radius': array([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]], dtype=float32)}
368class StructField(Field): 369 """Taichi struct field with SNode implementation. 370 371 Instead of directly constraining Expr entries, the StructField object 372 directly hosts members as `Field` instances to support nested structs. 373 374 Args: 375 field_dict (Dict[str, Field]): Struct field members. 376 struct_methods (Dict[str, callable]): Dictionary of functions to apply 377 to each struct instance in the field. 378 name (string, optional): The custom name of the field. 379 """ 380 381 def __init__(self, field_dict, struct_methods, name=None, is_primal=True): 382 # will not call Field initializer 383 self.field_dict = field_dict 384 self.struct_methods = struct_methods 385 self.name = name 386 self.grad = None 387 self.dual = None 388 if is_primal: 389 grad_field_dict = {} 390 for k, v in self.field_dict.items(): 391 grad_field_dict[k] = v.grad 392 self.grad = StructField(grad_field_dict, struct_methods, name + ".grad", is_primal=False) 393 394 dual_field_dict = {} 395 for k, v in self.field_dict.items(): 396 dual_field_dict[k] = v.dual 397 self.dual = StructField(dual_field_dict, struct_methods, name + ".dual", is_primal=False) 398 self._register_fields() 399 400 @property 401 def keys(self): 402 """Returns the list of names of the field members. 403 404 Example:: 405 406 >>> f1 = ti.Vector.field(3, ti.f32, shape=(3, 3)) 407 >>> f2 = ti.field(ti.f32, shape=(3, 3)) 408 >>> F = ti.StructField({"center": f1, "radius": f2}) 409 >>> F.keys 410 ['center', 'radius'] 411 """ 412 return list(self.field_dict.keys()) 413 414 @property 415 def _members(self): 416 return list(self.field_dict.values()) 417 418 @property 419 def _items(self): 420 return self.field_dict.items() 421 422 @staticmethod 423 def _make_getter(key): 424 def getter(self): 425 """Get an entry from custom struct by name.""" 426 return self.field_dict[key] 427 428 return getter 429 430 @staticmethod 431 def _make_setter(key): 432 @python_scope 433 def setter(self, value): 434 self.field_dict[key] = value 435 436 return setter 437 438 def _register_fields(self): 439 for k in self.keys: 440 setattr(self, k, self.field_dict[k]) 441 442 def _get_field_members(self): 443 """Gets A flattened list of all struct elements. 444 445 Returns: 446 A list of struct elements. 447 """ 448 field_members = [] 449 for m in self._members: 450 assert isinstance(m, Field) 451 field_members += m._get_field_members() 452 return field_members 453 454 @property 455 def _snode(self): 456 """Gets representative SNode for info purposes. 457 458 Returns: 459 SNode: Representative SNode (SNode of first field member). 460 """ 461 return self._members[0]._snode 462 463 def _loop_range(self): 464 """Gets SNode of representative field member for loop range info. 465 466 Returns: 467 taichi_python.SNode: SNode of representative (first) field member. 468 """ 469 return self._members[0]._loop_range() 470 471 @python_scope 472 def copy_from(self, other): 473 """Copies all elements from another field. 474 475 The shape of the other field needs to be the same as `self`. 476 477 Args: 478 other (Field): The source field. 479 """ 480 assert isinstance(other, Field) 481 assert set(self.keys) == set(other.keys) 482 for k in self.keys: 483 self.field_dict[k].copy_from(other.get_member_field(k)) 484 485 @python_scope 486 def fill(self, val): 487 """Fills this struct field with a specified value. 488 489 Args: 490 val (Union[int, float]): Value to fill. 491 """ 492 for v in self._members: 493 v.fill(val) 494 495 def _initialize_host_accessors(self): 496 for v in self._members: 497 v._initialize_host_accessors() 498 499 def get_member_field(self, key): 500 """Creates a ScalarField using a specific field member. 501 502 Args: 503 key (str): Specified key of the field member. 504 505 Returns: 506 ScalarField: The result ScalarField. 507 """ 508 return self.field_dict[key] 509 510 @python_scope 511 def from_numpy(self, array_dict): 512 """Copies the data from a set of `numpy.array` into this field. 513 514 The argument `array_dict` must be a dictionay-like object, it 515 contains all the keys in this field and the copying process 516 between corresponding items can be performed. 517 """ 518 for k, v in self._items: 519 v.from_numpy(array_dict[k]) 520 521 @python_scope 522 def from_torch(self, array_dict): 523 """Copies the data from a set of `torch.tensor` into this field. 524 525 The argument `array_dict` must be a dictionay-like object, it 526 contains all the keys in this field and the copying process 527 between corresponding items can be performed. 528 """ 529 for k, v in self._items: 530 v.from_torch(array_dict[k]) 531 532 @python_scope 533 def from_paddle(self, array_dict): 534 """Copies the data from a set of `paddle.Tensor` into this field. 535 536 The argument `array_dict` must be a dictionay-like object, it 537 contains all the keys in this field and the copying process 538 between corresponding items can be performed. 539 """ 540 for k, v in self._items: 541 v.from_paddle(array_dict[k]) 542 543 @python_scope 544 def to_numpy(self): 545 """Converts the Struct field instance to a dictionary of NumPy arrays. 546 547 The dictionary may be nested when converting nested structs. 548 549 Returns: 550 Dict[str, Union[numpy.ndarray, Dict]]: The result NumPy array. 551 """ 552 return {k: v.to_numpy() for k, v in self._items} 553 554 @python_scope 555 def to_torch(self, device=None): 556 """Converts the Struct field instance to a dictionary of PyTorch tensors. 557 558 The dictionary may be nested when converting nested structs. 559 560 Args: 561 device (torch.device, optional): The 562 desired device of returned tensor. 563 564 Returns: 565 Dict[str, Union[torch.Tensor, Dict]]: The result 566 PyTorch tensor. 567 """ 568 return {k: v.to_torch(device=device) for k, v in self._items} 569 570 @python_scope 571 def to_paddle(self, place=None): 572 """Converts the Struct field instance to a dictionary of Paddle tensors. 573 574 The dictionary may be nested when converting nested structs. 575 576 Args: 577 place (paddle.CPUPlace()/CUDAPlace(n), optional): The 578 desired place of returned tensor. 579 580 Returns: 581 Dict[str, Union[paddle.Tensor, Dict]]: The result 582 Paddle tensor. 583 """ 584 return {k: v.to_paddle(place=place) for k, v in self._items} 585 586 @python_scope 587 def __setitem__(self, indices, element): 588 self._initialize_host_accessors() 589 self[indices]._set_entries(element) 590 591 @python_scope 592 def __getitem__(self, indices): 593 self._initialize_host_accessors() 594 # scalar fields does not instantiate SNodeHostAccess by default 595 entries = { 596 k: v._host_access(self._pad_key(indices))[0] if isinstance(v, ScalarField) else v[indices] 597 for k, v in self._items 598 } 599 entries["__struct_methods"] = self.struct_methods 600 return Struct(entries)
Taichi struct field with SNode implementation.
Instead of directly constraining Expr entries, the StructField object
directly hosts members as Field instances to support nested structs.
Args: field_dict (Dict[str, Field]): Struct field members. struct_methods (Dict[str, callable]): Dictionary of functions to apply to each struct instance in the field. name (string, optional): The custom name of the field.
381 def __init__(self, field_dict, struct_methods, name=None, is_primal=True): 382 # will not call Field initializer 383 self.field_dict = field_dict 384 self.struct_methods = struct_methods 385 self.name = name 386 self.grad = None 387 self.dual = None 388 if is_primal: 389 grad_field_dict = {} 390 for k, v in self.field_dict.items(): 391 grad_field_dict[k] = v.grad 392 self.grad = StructField(grad_field_dict, struct_methods, name + ".grad", is_primal=False) 393 394 dual_field_dict = {} 395 for k, v in self.field_dict.items(): 396 dual_field_dict[k] = v.dual 397 self.dual = StructField(dual_field_dict, struct_methods, name + ".dual", is_primal=False) 398 self._register_fields()
400 @property 401 def keys(self): 402 """Returns the list of names of the field members. 403 404 Example:: 405 406 >>> f1 = ti.Vector.field(3, ti.f32, shape=(3, 3)) 407 >>> f2 = ti.field(ti.f32, shape=(3, 3)) 408 >>> F = ti.StructField({"center": f1, "radius": f2}) 409 >>> F.keys 410 ['center', 'radius'] 411 """ 412 return list(self.field_dict.keys())
Returns the list of names of the field members.
Example::
>>> f1 = ti.Vector.field(3, ti.f32, shape=(3, 3))
>>> f2 = ti.field(ti.f32, shape=(3, 3))
>>> F = ti.StructField({"center": f1, "radius": f2})
>>> F.keys
['center', 'radius']
471 @python_scope 472 def copy_from(self, other): 473 """Copies all elements from another field. 474 475 The shape of the other field needs to be the same as `self`. 476 477 Args: 478 other (Field): The source field. 479 """ 480 assert isinstance(other, Field) 481 assert set(self.keys) == set(other.keys) 482 for k in self.keys: 483 self.field_dict[k].copy_from(other.get_member_field(k))
Copies all elements from another field.
The shape of the other field needs to be the same as self.
Args: other (Field): The source field.
485 @python_scope 486 def fill(self, val): 487 """Fills this struct field with a specified value. 488 489 Args: 490 val (Union[int, float]): Value to fill. 491 """ 492 for v in self._members: 493 v.fill(val)
Fills this struct field with a specified value.
Args: val (Union[int, float]): Value to fill.
499 def get_member_field(self, key): 500 """Creates a ScalarField using a specific field member. 501 502 Args: 503 key (str): Specified key of the field member. 504 505 Returns: 506 ScalarField: The result ScalarField. 507 """ 508 return self.field_dict[key]
Creates a ScalarField using a specific field member.
Args: key (str): Specified key of the field member.
Returns: ScalarField: The result ScalarField.
510 @python_scope 511 def from_numpy(self, array_dict): 512 """Copies the data from a set of `numpy.array` into this field. 513 514 The argument `array_dict` must be a dictionay-like object, it 515 contains all the keys in this field and the copying process 516 between corresponding items can be performed. 517 """ 518 for k, v in self._items: 519 v.from_numpy(array_dict[k])
Copies the data from a set of numpy.array into this field.
The argument array_dict must be a dictionay-like object, it
contains all the keys in this field and the copying process
between corresponding items can be performed.
521 @python_scope 522 def from_torch(self, array_dict): 523 """Copies the data from a set of `torch.tensor` into this field. 524 525 The argument `array_dict` must be a dictionay-like object, it 526 contains all the keys in this field and the copying process 527 between corresponding items can be performed. 528 """ 529 for k, v in self._items: 530 v.from_torch(array_dict[k])
Copies the data from a set of torch.tensor into this field.
The argument array_dict must be a dictionay-like object, it
contains all the keys in this field and the copying process
between corresponding items can be performed.
532 @python_scope 533 def from_paddle(self, array_dict): 534 """Copies the data from a set of `paddle.Tensor` into this field. 535 536 The argument `array_dict` must be a dictionay-like object, it 537 contains all the keys in this field and the copying process 538 between corresponding items can be performed. 539 """ 540 for k, v in self._items: 541 v.from_paddle(array_dict[k])
Copies the data from a set of paddle.Tensor into this field.
The argument array_dict must be a dictionay-like object, it
contains all the keys in this field and the copying process
between corresponding items can be performed.
543 @python_scope 544 def to_numpy(self): 545 """Converts the Struct field instance to a dictionary of NumPy arrays. 546 547 The dictionary may be nested when converting nested structs. 548 549 Returns: 550 Dict[str, Union[numpy.ndarray, Dict]]: The result NumPy array. 551 """ 552 return {k: v.to_numpy() for k, v in self._items}
Converts the Struct field instance to a dictionary of NumPy arrays.
The dictionary may be nested when converting nested structs.
Returns: Dict[str, Union[numpy.ndarray, Dict]]: The result NumPy array.
554 @python_scope 555 def to_torch(self, device=None): 556 """Converts the Struct field instance to a dictionary of PyTorch tensors. 557 558 The dictionary may be nested when converting nested structs. 559 560 Args: 561 device (torch.device, optional): The 562 desired device of returned tensor. 563 564 Returns: 565 Dict[str, Union[torch.Tensor, Dict]]: The result 566 PyTorch tensor. 567 """ 568 return {k: v.to_torch(device=device) for k, v in self._items}
Converts the Struct field instance to a dictionary of PyTorch tensors.
The dictionary may be nested when converting nested structs.
Args: device (torch.device, optional): The desired device of returned tensor.
Returns: Dict[str, Union[torch.Tensor, Dict]]: The result PyTorch tensor.
570 @python_scope 571 def to_paddle(self, place=None): 572 """Converts the Struct field instance to a dictionary of Paddle tensors. 573 574 The dictionary may be nested when converting nested structs. 575 576 Args: 577 place (paddle.CPUPlace()/CUDAPlace(n), optional): The 578 desired place of returned tensor. 579 580 Returns: 581 Dict[str, Union[paddle.Tensor, Dict]]: The result 582 Paddle tensor. 583 """ 584 return {k: v.to_paddle(place=place) for k, v in self._items}
Converts the Struct field instance to a dictionary of Paddle tensors.
The dictionary may be nested when converting nested structs.
Args: place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor.
Returns: Dict[str, Union[paddle.Tensor, Dict]]: The result Paddle tensor.
43class TaichiAssertionError(TaichiRuntimeError, AssertionError): 44 """Thrown when assertion fails at runtime.""" 45 46 pass
Thrown when assertion fails at runtime.
7class TaichiCompilationError(Exception): 8 """Base class for all compilation exceptions.""" 9 10 pass
Base class for all compilation exceptions.
19class TaichiNameError(TaichiCompilationError, NameError): 20 """Thrown when an undefine name is found during compilation.""" 21 22 pass
Thrown when an undefine name is found during compilation.
37class TaichiRuntimeError(RuntimeError): 38 """Thrown when the compiled program cannot be executed due to unspecified reasons.""" 39 40 pass
Thrown when the compiled program cannot be executed due to unspecified reasons.
49class TaichiRuntimeTypeError(TaichiRuntimeError, TypeError): 50 @staticmethod 51 def get(pos, needed, provided): 52 return TaichiRuntimeTypeError( 53 f"Argument {pos} (type={provided}) cannot be converted into required type {needed}" 54 ) 55 56 @staticmethod 57 def get_ret(needed, provided): 58 return TaichiRuntimeTypeError(f"Return (type={provided}) cannot be converted into required type {needed}")
Thrown when the compiled program cannot be executed due to unspecified reasons.
13class TaichiSyntaxError(TaichiCompilationError, SyntaxError): 14 """Thrown when a syntax error is found during compilation.""" 15 16 pass
Thrown when a syntax error is found during compilation.
31class TaichiTypeError(TaichiCompilationError, TypeError): 32 """Thrown when a type mismatch is found during compilation.""" 33 34 pass
Thrown when a type mismatch is found during compilation.
101class Texture: 102 """Taichi Texture class. 103 104 Args: 105 fmt (ti.Format): Color format of the texture. 106 shape (Tuple[int]): Shape of the Texture. 107 """ 108 109 def __init__(self, fmt, arr_shape): 110 self.tex = impl.get_runtime().prog.create_texture(fmt, arr_shape) 111 self.fmt = fmt 112 self.num_dims = len(arr_shape) 113 self.shape = arr_shape 114 115 def from_ndarray(self, ndarray): 116 """Loads an ndarray to texture. 117 118 Args: 119 ndarray (ti.Ndarray): Source ndarray to load from. 120 """ 121 self.tex.from_ndarray(ndarray.arr) 122 123 def from_field(self, field): 124 """Loads a field to texture. 125 126 Args: 127 field (ti.Field): Source field to load from. 128 """ 129 self.tex.from_snode(field.snode.ptr) 130 131 def _device_allocation_ptr(self): 132 return self.tex.device_allocation_ptr() 133 134 def from_image(self, image): 135 """Loads a PIL image to texture. This method is only allowed a 2D texture with `ti.Format.rgba8`. 136 137 Args: 138 image (PIL.Image.Image): Source PIL image to load from. 139 140 """ 141 from PIL import Image # pylint: disable=import-outside-toplevel 142 143 assert isinstance(image, Image.Image) 144 if image.mode != "RGB": 145 image = image.convert("RGB") 146 assert image.size == tuple(self.shape) 147 148 assert self.num_dims == 2 149 # Don't use transpose method since its enums are too new 150 image = image.rotate(90, expand=True) 151 arr = np.asarray(image) 152 from taichi._kernels import ( # pylint: disable=import-outside-toplevel 153 load_texture_from_numpy, 154 ) 155 156 load_texture_from_numpy(self, arr) 157 158 def to_image(self): 159 """Saves a texture to a PIL image in RGB mode. This method is only allowed a 2D texture with `ti.Format.rgba8`. 160 161 Returns: 162 img (PIL.Image.Image): a PIL image in RGB mode, with the same size as source texture. 163 """ 164 assert self.num_dims == 2 165 from PIL import Image # pylint: disable=import-outside-toplevel 166 167 res = np.zeros(self.shape + (3,), np.uint8) 168 from taichi._kernels import ( # pylint: disable=import-outside-toplevel 169 save_texture_to_numpy, 170 ) 171 172 save_texture_to_numpy(self, res) 173 return Image.fromarray(res).rotate(270, expand=True)
Taichi Texture class.
Args: fmt (ti.Format): Color format of the texture. shape (Tuple[int]): Shape of the Texture.
115 def from_ndarray(self, ndarray): 116 """Loads an ndarray to texture. 117 118 Args: 119 ndarray (ti.Ndarray): Source ndarray to load from. 120 """ 121 self.tex.from_ndarray(ndarray.arr)
Loads an ndarray to texture.
Args: ndarray (ti.Ndarray): Source ndarray to load from.
123 def from_field(self, field): 124 """Loads a field to texture. 125 126 Args: 127 field (ti.Field): Source field to load from. 128 """ 129 self.tex.from_snode(field.snode.ptr)
Loads a field to texture.
Args: field (ti.Field): Source field to load from.
134 def from_image(self, image): 135 """Loads a PIL image to texture. This method is only allowed a 2D texture with `ti.Format.rgba8`. 136 137 Args: 138 image (PIL.Image.Image): Source PIL image to load from. 139 140 """ 141 from PIL import Image # pylint: disable=import-outside-toplevel 142 143 assert isinstance(image, Image.Image) 144 if image.mode != "RGB": 145 image = image.convert("RGB") 146 assert image.size == tuple(self.shape) 147 148 assert self.num_dims == 2 149 # Don't use transpose method since its enums are too new 150 image = image.rotate(90, expand=True) 151 arr = np.asarray(image) 152 from taichi._kernels import ( # pylint: disable=import-outside-toplevel 153 load_texture_from_numpy, 154 ) 155 156 load_texture_from_numpy(self, arr)
Loads a PIL image to texture. This method is only allowed a 2D texture with ti.Format.rgba8.
Args: image (PIL.Image.Image): Source PIL image to load from.
158 def to_image(self): 159 """Saves a texture to a PIL image in RGB mode. This method is only allowed a 2D texture with `ti.Format.rgba8`. 160 161 Returns: 162 img (PIL.Image.Image): a PIL image in RGB mode, with the same size as source texture. 163 """ 164 assert self.num_dims == 2 165 from PIL import Image # pylint: disable=import-outside-toplevel 166 167 res = np.zeros(self.shape + (3,), np.uint8) 168 from taichi._kernels import ( # pylint: disable=import-outside-toplevel 169 save_texture_to_numpy, 170 ) 171 172 save_texture_to_numpy(self, res) 173 return Image.fromarray(res).rotate(270, expand=True)
Saves a texture to a PIL image in RGB mode. This method is only allowed a 2D texture with ti.Format.rgba8.
Returns: img (PIL.Image.Image): a PIL image in RGB mode, with the same size as source texture.
1093class Vector(Matrix): 1094 def __init__(self, arr, dt=None, **kwargs): 1095 """Constructs a vector from given array. 1096 1097 A vector is an instance of a 2-D matrix with the second dimension being equal to 1. 1098 1099 Args: 1100 arr (Union[list, tuple, np.ndarray]): The initial values of the Vector. 1101 dt (:mod:`~taichi.types.primitive_types`): data type of the vector. 1102 1103 Returns: 1104 :class:`~taichi.Matrix`: A vector instance. 1105 Example:: 1106 >>> u = ti.Vector([1, 2]) 1107 >>> print(u.m, u.n) # verify a vector is a matrix of shape (n, 1) 1108 2 1 1109 >>> v = ti.Vector([3, 4]) 1110 >>> u + v 1111 [4 6] 1112 """ 1113 super().__init__(arr, dt=dt, **kwargs) 1114 1115 def get_shape(self): 1116 return (self.n,) 1117 1118 @classmethod 1119 def field(cls, n, dtype, *args, **kwargs): 1120 """ti.Vector.field""" 1121 ndim = kwargs.get("ndim", 1) 1122 assert ndim == 1 1123 kwargs["ndim"] = 1 1124 return super().field(n, 1, dtype, *args, **kwargs) 1125 1126 @classmethod 1127 @python_scope 1128 def ndarray(cls, n, dtype, shape): 1129 """Defines a Taichi ndarray with vector elements. 1130 1131 Args: 1132 n (int): Size of the vector. 1133 dtype (DataType): Data type of each value. 1134 shape (Union[int, tuple[int]]): Shape of the ndarray. 1135 1136 Example: 1137 The code below shows how a Taichi ndarray with vector elements can be declared and defined:: 1138 1139 >>> x = ti.Vector.ndarray(3, ti.f32, shape=(16, 8)) 1140 """ 1141 if isinstance(shape, numbers.Number): 1142 shape = (shape,) 1143 return VectorNdarray(n, dtype, shape)
The matrix class.
A matrix is a 2-D rectangular array with scalar entries, it's row-majored, and is aligned continuously. We recommend only use matrix with no more than 32 elements for efficiency considerations.
Note: in taichi a matrix is strictly two-dimensional and only stores scalars.
Args:
arr (Union[list, tuple, np.ndarray]): the initial values of a matrix.
dt (~taichi.types.primitive_types): the element data type.
ndim (int optional): the number of dimensions of the matrix; forced reshape if given.
Example::
use a 2d list to initialize a matrix
>>> @ti.kernel
>>> def test():
>>> n = 5
>>> M = ti.Matrix([[0] * n for _ in range(n)], ti.i32)
>>> print(M) # a 5x5 matrix with integer elements
get the number of rows and columns via the `n`, `m` property:
>>> M = ti.Matrix([[0, 1], [2, 3], [4, 5]], ti.i32)
>>> M.n # number of rows
3
>>> M.m # number of cols
>>> 2
you can even initialize a matrix with an empty list:
>>> M = ti.Matrix([[], []], ti.i32)
>>> M.n
2
>>> M.m
0
1094 def __init__(self, arr, dt=None, **kwargs): 1095 """Constructs a vector from given array. 1096 1097 A vector is an instance of a 2-D matrix with the second dimension being equal to 1. 1098 1099 Args: 1100 arr (Union[list, tuple, np.ndarray]): The initial values of the Vector. 1101 dt (:mod:`~taichi.types.primitive_types`): data type of the vector. 1102 1103 Returns: 1104 :class:`~taichi.Matrix`: A vector instance. 1105 Example:: 1106 >>> u = ti.Vector([1, 2]) 1107 >>> print(u.m, u.n) # verify a vector is a matrix of shape (n, 1) 1108 2 1 1109 >>> v = ti.Vector([3, 4]) 1110 >>> u + v 1111 [4 6] 1112 """ 1113 super().__init__(arr, dt=dt, **kwargs)
Constructs a vector from given array.
A vector is an instance of a 2-D matrix with the second dimension being equal to 1.
Args:
arr (Union[list, tuple, np.ndarray]): The initial values of the Vector.
dt (~taichi.types.primitive_types): data type of the vector.
Returns:
~taichi.Matrix: A vector instance.
Example::
u = ti.Vector([1, 2]) print(u.m, u.n) # verify a vector is a matrix of shape (n, 1) 2 1 v = ti.Vector([3, 4]) u + v [4 6]
1118 @classmethod 1119 def field(cls, n, dtype, *args, **kwargs): 1120 """ti.Vector.field""" 1121 ndim = kwargs.get("ndim", 1) 1122 assert ndim == 1 1123 kwargs["ndim"] = 1 1124 return super().field(n, 1, dtype, *args, **kwargs)
ti.Vector.field
1126 @classmethod 1127 @python_scope 1128 def ndarray(cls, n, dtype, shape): 1129 """Defines a Taichi ndarray with vector elements. 1130 1131 Args: 1132 n (int): Size of the vector. 1133 dtype (DataType): Data type of each value. 1134 shape (Union[int, tuple[int]]): Shape of the ndarray. 1135 1136 Example: 1137 The code below shows how a Taichi ndarray with vector elements can be declared and defined:: 1138 1139 >>> x = ti.Vector.ndarray(3, ti.f32, shape=(16, 8)) 1140 """ 1141 if isinstance(shape, numbers.Number): 1142 shape = (shape,) 1143 return VectorNdarray(n, dtype, shape)
Defines a Taichi ndarray with vector elements.
Args: n (int): Size of the vector. dtype (DataType): Data type of each value. shape (Union[int, tuple[int]]): Shape of the ndarray.
Example: The code below shows how a Taichi ndarray with vector elements can be declared and defined::
>>> x = ti.Vector.ndarray(3, ti.f32, shape=(16, 8))
1769class VectorNdarray(Ndarray): 1770 """Taichi ndarray with vector elements. 1771 1772 Args: 1773 n (int): Size of the vector. 1774 dtype (DataType): Data type of each value. 1775 shape (Tuple[int]): Shape of the ndarray. 1776 1777 Example:: 1778 1779 >>> a = ti.VectorNdarray(3, ti.f32, (3, 3)) 1780 """ 1781 1782 def __init__(self, n, dtype, shape): 1783 self.n = n 1784 super().__init__() 1785 # TODO(zhanlue): remove self.dtype and migrate its usages to element_type 1786 self.dtype = cook_dtype(dtype) 1787 1788 self.layout = Layout.AOS 1789 self.shape = tuple(shape) 1790 self.element_type = _type_factory.get_tensor_type((n,), self.dtype) 1791 self.arr = impl.get_runtime().prog.create_ndarray( 1792 cook_dtype(self.element_type), 1793 shape, 1794 Layout.AOS, 1795 zero_fill=True, 1796 dbg_info=ti_python_core.DebugInfo(get_traceback()), 1797 ) 1798 1799 @property 1800 def element_shape(self): 1801 """Gets the dimension of the vector of this ndarray. 1802 1803 Example:: 1804 1805 >>> a = ti.VectorNdarray(3, ti.f32, (3, 3)) 1806 >>> a.element_shape 1807 (3,) 1808 """ 1809 return tuple(self.arr.element_shape()) 1810 1811 @python_scope 1812 def __setitem__(self, key, value): 1813 if not isinstance(value, (list, tuple)): 1814 value = list(value) 1815 for i in range(self.n): 1816 self[key][i] = value[i] 1817 1818 @python_scope 1819 def __getitem__(self, key): 1820 key = () if key is None else (key,) if isinstance(key, numbers.Number) else tuple(key) 1821 return Vector([NdarrayHostAccess(self, key, (i,)) for i in range(self.n)]) 1822 1823 @python_scope 1824 def to_numpy(self): 1825 """Converts this vector ndarray to a `numpy.ndarray`. 1826 1827 Example:: 1828 1829 >>> a = ti.VectorNdarray(3, ti.f32, (2, 2)) 1830 >>> a.to_numpy() 1831 array([[[0., 0., 0.], 1832 [0., 0., 0.]], 1833 1834 [[0., 0., 0.], 1835 [0., 0., 0.]]], dtype=float32) 1836 """ 1837 return self._ndarray_matrix_to_numpy(as_vector=1) 1838 1839 @python_scope 1840 def from_numpy(self, arr): 1841 """Copies the data from a `numpy.ndarray` into this ndarray. 1842 1843 The shape and data type of `arr` must match this ndarray. 1844 1845 Example:: 1846 1847 >>> import numpy as np 1848 >>> a = ti.VectorNdarray(3, ti.f32, (2, 2), 0) 1849 >>> b = np.ones((2, 2, 3), dtype=np.float32) 1850 >>> a.from_numpy(b) 1851 """ 1852 self._ndarray_matrix_from_numpy(arr, as_vector=1) 1853 1854 @python_scope 1855 def __deepcopy__(self, memo=None): 1856 ret_arr = VectorNdarray(self.n, self.dtype, self.shape) 1857 ret_arr.copy_from(self) 1858 return ret_arr 1859 1860 @python_scope 1861 def _fill_by_kernel(self, val): 1862 from taichi._kernels import fill_ndarray_matrix # pylint: disable=C0415 1863 1864 shape = self.element_type.shape() 1865 prim_dtype = self.element_type.element_type() 1866 vector_type = VectorType(shape[0], prim_dtype) 1867 if isinstance(val, Vector): 1868 value = val 1869 else: 1870 value = vector_type(val) 1871 fill_ndarray_matrix(self, value) 1872 1873 @python_scope 1874 def __repr__(self): 1875 return f"<{self.n} {Layout.AOS} ti.Vector.ndarray>"
Taichi ndarray with vector elements.
Args: n (int): Size of the vector. dtype (DataType): Data type of each value. shape (Tuple[int]): Shape of the ndarray.
Example::
>>> a = ti.VectorNdarray(3, ti.f32, (3, 3))
1782 def __init__(self, n, dtype, shape): 1783 self.n = n 1784 super().__init__() 1785 # TODO(zhanlue): remove self.dtype and migrate its usages to element_type 1786 self.dtype = cook_dtype(dtype) 1787 1788 self.layout = Layout.AOS 1789 self.shape = tuple(shape) 1790 self.element_type = _type_factory.get_tensor_type((n,), self.dtype) 1791 self.arr = impl.get_runtime().prog.create_ndarray( 1792 cook_dtype(self.element_type), 1793 shape, 1794 Layout.AOS, 1795 zero_fill=True, 1796 dbg_info=ti_python_core.DebugInfo(get_traceback()), 1797 )
1799 @property 1800 def element_shape(self): 1801 """Gets the dimension of the vector of this ndarray. 1802 1803 Example:: 1804 1805 >>> a = ti.VectorNdarray(3, ti.f32, (3, 3)) 1806 >>> a.element_shape 1807 (3,) 1808 """ 1809 return tuple(self.arr.element_shape())
Gets the dimension of the vector of this ndarray.
Example::
>>> a = ti.VectorNdarray(3, ti.f32, (3, 3))
>>> a.element_shape
(3,)
1823 @python_scope 1824 def to_numpy(self): 1825 """Converts this vector ndarray to a `numpy.ndarray`. 1826 1827 Example:: 1828 1829 >>> a = ti.VectorNdarray(3, ti.f32, (2, 2)) 1830 >>> a.to_numpy() 1831 array([[[0., 0., 0.], 1832 [0., 0., 0.]], 1833 1834 [[0., 0., 0.], 1835 [0., 0., 0.]]], dtype=float32) 1836 """ 1837 return self._ndarray_matrix_to_numpy(as_vector=1)
Converts this vector ndarray to a numpy.ndarray.
Example::
>>> a = ti.VectorNdarray(3, ti.f32, (2, 2))
>>> a.to_numpy()
array([[[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.]]], dtype=float32)
1839 @python_scope 1840 def from_numpy(self, arr): 1841 """Copies the data from a `numpy.ndarray` into this ndarray. 1842 1843 The shape and data type of `arr` must match this ndarray. 1844 1845 Example:: 1846 1847 >>> import numpy as np 1848 >>> a = ti.VectorNdarray(3, ti.f32, (2, 2), 0) 1849 >>> b = np.ones((2, 2, 3), dtype=np.float32) 1850 >>> a.from_numpy(b) 1851 """ 1852 self._ndarray_matrix_from_numpy(arr, as_vector=1)
Copies the data from a numpy.ndarray into this ndarray.
The shape and data type of arr must match this ndarray.
Example::
>>> import numpy as np
>>> a = ti.VectorNdarray(3, ti.f32, (2, 2), 0)
>>> b = np.ones((2, 2, 3), dtype=np.float32)
>>> a.from_numpy(b)
508def abs(x): # pylint: disable=W0622 509 """Compute the absolute value :math:`|x|` of `x`, element-wise. 510 511 Args: 512 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 513 Input scalar or matrix. 514 515 Returns: 516 The absolute value of each element in `x`. 517 518 Example:: 519 520 >>> @ti.kernel 521 >>> def test(): 522 >>> x = ti.Vector([-1.0, 0.0, 1.0]) 523 >>> y = ti.abs(x) 524 >>> print(y) 525 >>> 526 >>> test() 527 [1.0, 0.0, 1.0] 528 """ 529 return _unary_operation(_ti_core.expr_abs, builtins.abs, x)
Compute the absolute value \( |x| \) of x, element-wise.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
Returns:
The absolute value of each element in x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([-1.0, 0.0, 1.0])
>>> y = ti.abs(x)
>>> print(y)
>>>
>>> test()
[1.0, 0.0, 1.0]
244def acos(x): 245 """Trigonometric inverse cosine, element-wise. 246 247 The inverse of `cos` so that, if `y = cos(x)`, then `x = acos(y)`. 248 249 For input `x` not in the domain `[-1, 1]`, this function returns `nan` if \ 250 it's called in taichi scope, or raises exception if it's called in python scope. 251 252 Args: 253 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 254 A scalar or a matrix with elements in [-1, 1]. 255 256 Returns: 257 The inverse cosine of each element in `x`, in radians and in the closed \ 258 interval `[0, pi]`. This is a scalar if `x` is a scalar. 259 260 Example:: 261 262 >>> from math import pi 263 >>> ti.acos(ti.Matrix([-1.0, 0.0, 1.0])) * 180 / pi 264 [180., 90., 0.] 265 """ 266 return _unary_operation(_ti_core.expr_acos, np.arccos, x)
Trigonometric inverse cosine, element-wise.
The inverse of cos so that, if y = cos(x), then x = acos(y).
For input x not in the domain [-1, 1], this function returns nan if it's called in taichi scope, or raises exception if it's called in python scope.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): A scalar or a matrix with elements in [-1, 1].
Returns:
The inverse cosine of each element in x, in radians and in the closed interval [0, pi]. This is a scalar if x is a scalar.
Example::
>>> from math import pi
>>> ti.acos(ti.Matrix([-1.0, 0.0, 1.0])) * 180 / pi
[180., 90., 0.]
414def activate(node, indices): 415 """Explicitly activate a cell of `node` at location `indices`. 416 417 Args: 418 node (:class:`~taichi.SNode`): Must be a pointer, hash or bitmasked node. 419 indices (Union[int, :class:`~taichi.Vector`]): the indices to activate. 420 """ 421 impl.get_runtime().compiling_callable.ast_builder().insert_activate( 422 node._snode.ptr, expr.make_expr_group(indices), _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()) 423 )
Explicitly activate a cell of node at location indices.
Args:
node (~taichi.SNode): Must be a pointer, hash or bitmasked node.
indices (Union[int, ~taichi.Vector]): the indices to activate.
376def append(node, indices, val): 377 """Append a value `val` to a SNode `node` at index `indices`. 378 379 Args: 380 node (:class:`~taichi.SNode`): Input SNode. 381 indices (Union[int, :class:`~taichi.Vector`]): the indices to visit. 382 val (Union[:mod:`~taichi.types.primitive_types`, :mod:`~taichi.types.compound_types`]): the data to be appended. 383 """ 384 ptrs = expr._get_flattened_ptrs(val) 385 append_expr = expr.Expr( 386 impl.get_runtime() 387 .compiling_callable.ast_builder() 388 .expr_snode_append(node._snode.ptr, expr.make_expr_group(indices), ptrs), 389 dbg_info=_ti_core.DebugInfo(impl.get_runtime().get_current_src_info()), 390 ) 391 a = impl.expr_init(append_expr) 392 return a
Append a value val to a SNode node at index indices.
Args:
node (~taichi.SNode): Input SNode.
indices (Union[int, ~taichi.Vector]): the indices to visit.
val (Union[~taichi.types.primitive_types, ~taichi.types.compound_types]): the data to be appended.
219def asin(x): 220 """Trigonometric inverse sine, element-wise. 221 222 The inverse of `sin` so that, if `y = sin(x)`, then `x = asin(y)`. 223 224 For input `x` not in the domain `[-1, 1]`, this function returns `nan` if \ 225 it's called in taichi scope, or raises exception if it's called in python scope. 226 227 Args: 228 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 229 A scalar or a matrix with elements in [-1, 1]. 230 231 Returns: 232 The inverse sine of each element in `x`, in radians and in the closed \ 233 interval `[-pi/2, pi/2]`. 234 235 Example:: 236 237 >>> from math import pi 238 >>> ti.asin(ti.Matrix([-1.0, 0.0, 1.0])) * 180 / pi 239 [-90., 0., 90.] 240 """ 241 return _unary_operation(_ti_core.expr_asin, np.arcsin, x)
Trigonometric inverse sine, element-wise.
The inverse of sin so that, if y = sin(x), then x = asin(y).
For input x not in the domain [-1, 1], this function returns nan if it's called in taichi scope, or raises exception if it's called in python scope.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): A scalar or a matrix with elements in [-1, 1].
Returns:
The inverse sine of each element in x, in radians and in the closed interval [-pi/2, pi/2].
Example::
>>> from math import pi
>>> ti.asin(ti.Matrix([-1.0, 0.0, 1.0])) * 180 / pi
[-90., 0., 90.]
544def assume_in_range(val, base, low, high): 545 """Hints the compiler that a value is between a specified range, 546 for the compiler to perform scatchpad optimization, and return the 547 value untouched. 548 549 The assumed range is `[base + low, base + high)`. 550 551 Args: 552 553 val (Number): The input value. 554 base (Number): The base point for the range interval. 555 low (Number): The lower offset relative to `base` (included). 556 high (Number): The higher offset relative to `base` (excluded). 557 558 Returns: 559 Return the input `value` untouched. 560 561 Example:: 562 563 >>> # hint the compiler that x is in range [8, 12). 564 >>> x = ti.assume_in_range(x, 10, -2, 2) 565 >>> x 566 10 567 """ 568 return _ti_core.expr_assume_in_range( 569 Expr(val).ptr, Expr(base).ptr, low, high, _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()) 570 )
Hints the compiler that a value is between a specified range, for the compiler to perform scatchpad optimization, and return the value untouched.
The assumed range is [base + low, base + high).
Args:
val (Number): The input value.
base (Number): The base point for the range interval.
low (Number): The lower offset relative to `base` (included).
high (Number): The higher offset relative to `base` (excluded).
Returns:
Return the input value untouched.
Example::
>>> # hint the compiler that x is in range [8, 12).
>>> x = ti.assume_in_range(x, 10, -2, 2)
>>> x
10
775def atan2(x1, x2): 776 """Element-wise arc tangent of `x1/x2`. 777 778 Args: 779 x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 780 y-coordinates. 781 x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 782 x-coordinates. 783 784 Returns: 785 Angles in radians, in the range `[-pi, pi]`. 786 This is a scalar if both `x1` and `x2` are scalars. 787 788 Example:: 789 790 >>> from math import pi 791 >>> @ti.kernel 792 >>> def test(): 793 >>> x = ti.Matrix([-1.0, 1.0, -1.0, 1.0]) 794 >>> y = ti.Matrix([-1.0, -1.0, 1.0, 1.0]) 795 >>> z = ti.atan2(y, x) * 180 / pi 796 >>> print(z) 797 >>> 798 >>> test() 799 [-135.0, -45.0, 135.0, 45.0] 800 """ 801 return _binary_operation(_ti_core.expr_atan2, np.arctan2, x1, x2)
Element-wise arc tangent of x1/x2.
Args:
x1 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): y-coordinates.
x2 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): x-coordinates.
Returns:
Angles in radians, in the range [-pi, pi].
This is a scalar if both x1 and x2 are scalars.
Example::
>>> from math import pi
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([-1.0, 1.0, -1.0, 1.0])
>>> y = ti.Matrix([-1.0, -1.0, 1.0, 1.0])
>>> z = ti.atan2(y, x) * 180 / pi
>>> print(z)
>>>
>>> test()
[-135.0, -45.0, 135.0, 45.0]
1139@writeback_binary 1140def atomic_add(x, y): 1141 """Atomically compute `x + y`, store the result in `x`, 1142 and return the old value of `x`. 1143 1144 `x` must be a writable target, constant expressions or scalars 1145 are not allowed. 1146 1147 Args: 1148 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1149 The input. 1150 1151 Returns: 1152 The old value of `x`. 1153 1154 Example:: 1155 1156 >>> @ti.kernel 1157 >>> def test(): 1158 >>> x = ti.Vector([0, 0, 0]) 1159 >>> y = ti.Vector([1, 2, 3]) 1160 >>> z = ti.atomic_add(x, y) 1161 >>> print(x) # [1, 2, 3] the new value of x 1162 >>> print(z) # [0, 0, 0], the old value of x 1163 >>> 1164 >>> ti.atomic_add(1, x) # will raise TaichiSyntaxError 1165 """ 1166 return impl.expr_init(expr.Expr(_ti_core.expr_atomic_add(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())))
Atomically compute x + y, store the result in x,
and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([0, 0, 0])
>>> y = ti.Vector([1, 2, 3])
>>> z = ti.atomic_add(x, y)
>>> print(x) # [1, 2, 3] the new value of x
>>> print(z) # [0, 0, 0], the old value of x
>>>
>>> ti.atomic_add(1, x) # will raise TaichiSyntaxError
1289@writeback_binary 1290def atomic_and(x, y): 1291 """Atomically compute the bit-wise AND of `x` and `y`, element-wise. 1292 Store the result in `x`, and return the old value of `x`. 1293 1294 `x` must be a writable target, constant expressions or scalars 1295 are not allowed. 1296 1297 Args: 1298 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1299 The input. When both are matrices they must have the same shape. 1300 1301 Returns: 1302 The old value of `x`. 1303 1304 Example:: 1305 1306 >>> @ti.kernel 1307 >>> def test(): 1308 >>> x = ti.Vector([-1, 0, 1]) 1309 >>> y = ti.Vector([1, 2, 3]) 1310 >>> z = ti.atomic_and(x, y) 1311 >>> print(x) # [1, 0, 1] the new value of x 1312 >>> print(z) # [-1, 0, 1], the old value of x 1313 >>> 1314 >>> ti.atomic_and(1, x) # will raise TaichiSyntaxError 1315 """ 1316 return impl.expr_init( 1317 expr.Expr(_ti_core.expr_atomic_bit_and(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())) 1318 )
Atomically compute the bit-wise AND of x and y, element-wise.
Store the result in x, and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input. When both are matrices they must have the same shape.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([-1, 0, 1])
>>> y = ti.Vector([1, 2, 3])
>>> z = ti.atomic_and(x, y)
>>> print(x) # [1, 0, 1] the new value of x
>>> print(z) # [-1, 0, 1], the old value of x
>>>
>>> ti.atomic_and(1, x) # will raise TaichiSyntaxError
1259@writeback_binary 1260def atomic_max(x, y): 1261 """Atomically compute the maximum of `x` and `y`, element-wise. 1262 Store the result in `x`, and return the old value of `x`. 1263 1264 `x` must be a writable target, constant expressions or scalars 1265 are not allowed. 1266 1267 Args: 1268 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1269 The input. 1270 1271 Returns: 1272 The old value of `x`. 1273 1274 Example:: 1275 1276 >>> @ti.kernel 1277 >>> def test(): 1278 >>> x = 1 1279 >>> y = 2 1280 >>> z = ti.atomic_max(x, y) 1281 >>> print(x) # 2 the new value of x 1282 >>> print(z) # 1, the old value of x 1283 >>> 1284 >>> ti.atomic_max(1, x) # will raise TaichiSyntaxError 1285 """ 1286 return impl.expr_init(expr.Expr(_ti_core.expr_atomic_max(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())))
Atomically compute the maximum of x and y, element-wise.
Store the result in x, and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = 1
>>> y = 2
>>> z = ti.atomic_max(x, y)
>>> print(x) # 2 the new value of x
>>> print(z) # 1, the old value of x
>>>
>>> ti.atomic_max(1, x) # will raise TaichiSyntaxError
1229@writeback_binary 1230def atomic_min(x, y): 1231 """Atomically compute the minimum of `x` and `y`, element-wise. 1232 Store the result in `x`, and return the old value of `x`. 1233 1234 `x` must be a writable target, constant expressions or scalars 1235 are not allowed. 1236 1237 Args: 1238 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1239 The input. 1240 1241 Returns: 1242 The old value of `x`. 1243 1244 Example:: 1245 1246 >>> @ti.kernel 1247 >>> def test(): 1248 >>> x = 2 1249 >>> y = 1 1250 >>> z = ti.atomic_min(x, y) 1251 >>> print(x) # 1 the new value of x 1252 >>> print(z) # 2, the old value of x 1253 >>> 1254 >>> ti.atomic_min(1, x) # will raise TaichiSyntaxError 1255 """ 1256 return impl.expr_init(expr.Expr(_ti_core.expr_atomic_min(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())))
Atomically compute the minimum of x and y, element-wise.
Store the result in x, and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = 2
>>> y = 1
>>> z = ti.atomic_min(x, y)
>>> print(x) # 1 the new value of x
>>> print(z) # 2, the old value of x
>>>
>>> ti.atomic_min(1, x) # will raise TaichiSyntaxError
1169@writeback_binary 1170def atomic_mul(x, y): 1171 """Atomically compute `x * y`, store the result in `x`, 1172 and return the old value of `x`. 1173 1174 `x` must be a writable target, constant expressions or scalars 1175 are not allowed. 1176 1177 Args: 1178 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1179 The input. 1180 1181 Returns: 1182 The old value of `x`. 1183 1184 Example:: 1185 1186 >>> @ti.kernel 1187 >>> def test(): 1188 >>> x = ti.Vector([1, 2, 3]) 1189 >>> y = ti.Vector([4, 5, 6]) 1190 >>> z = ti.atomic_mul(x, y) 1191 >>> print(x) # [1, 2, 3] the new value of x 1192 >>> print(z) # [4, 10, 18], the old value of x 1193 >>> 1194 >>> ti.atomic_mul(1, x) # will raise TaichiSyntaxError 1195 """ 1196 return impl.expr_init(expr.Expr(_ti_core.expr_atomic_mul(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())))
Atomically compute x * y, store the result in x,
and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([1, 2, 3])
>>> y = ti.Vector([4, 5, 6])
>>> z = ti.atomic_mul(x, y)
>>> print(x) # [1, 2, 3] the new value of x
>>> print(z) # [4, 10, 18], the old value of x
>>>
>>> ti.atomic_mul(1, x) # will raise TaichiSyntaxError
1321@writeback_binary 1322def atomic_or(x, y): 1323 """Atomically compute the bit-wise OR of `x` and `y`, element-wise. 1324 Store the result in `x`, and return the old value of `x`. 1325 1326 `x` must be a writable target, constant expressions or scalars 1327 are not allowed. 1328 1329 Args: 1330 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1331 The input. When both are matrices they must have the same shape. 1332 1333 Returns: 1334 The old value of `x`. 1335 1336 Example:: 1337 1338 >>> @ti.kernel 1339 >>> def test(): 1340 >>> x = ti.Vector([-1, 0, 1]) 1341 >>> y = ti.Vector([1, 2, 3]) 1342 >>> z = ti.atomic_or(x, y) 1343 >>> print(x) # [-1, 2, 3] the new value of x 1344 >>> print(z) # [-1, 0, 1], the old value of x 1345 >>> 1346 >>> ti.atomic_or(1, x) # will raise TaichiSyntaxError 1347 """ 1348 return impl.expr_init( 1349 expr.Expr(_ti_core.expr_atomic_bit_or(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())) 1350 )
Atomically compute the bit-wise OR of x and y, element-wise.
Store the result in x, and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input. When both are matrices they must have the same shape.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([-1, 0, 1])
>>> y = ti.Vector([1, 2, 3])
>>> z = ti.atomic_or(x, y)
>>> print(x) # [-1, 2, 3] the new value of x
>>> print(z) # [-1, 0, 1], the old value of x
>>>
>>> ti.atomic_or(1, x) # will raise TaichiSyntaxError
1199@writeback_binary 1200def atomic_sub(x, y): 1201 """Atomically subtract `x` by `y`, store the result in `x`, 1202 and return the old value of `x`. 1203 1204 `x` must be a writable target, constant expressions or scalars 1205 are not allowed. 1206 1207 Args: 1208 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1209 The input. 1210 1211 Returns: 1212 The old value of `x`. 1213 1214 Example:: 1215 1216 >>> @ti.kernel 1217 >>> def test(): 1218 >>> x = ti.Vector([0, 0, 0]) 1219 >>> y = ti.Vector([1, 2, 3]) 1220 >>> z = ti.atomic_sub(x, y) 1221 >>> print(x) # [-1, -2, -3] the new value of x 1222 >>> print(z) # [0, 0, 0], the old value of x 1223 >>> 1224 >>> ti.atomic_sub(1, x) # will raise TaichiSyntaxError 1225 """ 1226 return impl.expr_init(expr.Expr(_ti_core.expr_atomic_sub(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())))
Atomically subtract x by y, store the result in x,
and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([0, 0, 0])
>>> y = ti.Vector([1, 2, 3])
>>> z = ti.atomic_sub(x, y)
>>> print(x) # [-1, -2, -3] the new value of x
>>> print(z) # [0, 0, 0], the old value of x
>>>
>>> ti.atomic_sub(1, x) # will raise TaichiSyntaxError
1353@writeback_binary 1354def atomic_xor(x, y): 1355 """Atomically compute the bit-wise XOR of `x` and `y`, element-wise. 1356 Store the result in `x`, and return the old value of `x`. 1357 1358 `x` must be a writable target, constant expressions or scalars 1359 are not allowed. 1360 1361 Args: 1362 x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1363 The input. When both are matrices they must have the same shape. 1364 1365 Returns: 1366 The old value of `x`. 1367 1368 Example:: 1369 1370 >>> @ti.kernel 1371 >>> def test(): 1372 >>> x = ti.Vector([-1, 0, 1]) 1373 >>> y = ti.Vector([1, 2, 3]) 1374 >>> z = ti.atomic_xor(x, y) 1375 >>> print(x) # [-2, 2, 2] the new value of x 1376 >>> print(z) # [-1, 0, 1], the old value of x 1377 >>> 1378 >>> ti.atomic_xor(1, x) # will raise TaichiSyntaxError 1379 """ 1380 return impl.expr_init( 1381 expr.Expr(_ti_core.expr_atomic_bit_xor(x.ptr, y.ptr), dbg_info=_ti_core.DebugInfo(stack_info())) 1382 )
Atomically compute the bit-wise XOR of x and y, element-wise.
Store the result in x, and return the old value of x.
x must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input. When both are matrices they must have the same shape.
Returns:
The old value of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([-1, 0, 1])
>>> y = ti.Vector([1, 2, 3])
>>> z = ti.atomic_xor(x, y)
>>> print(x) # [-2, 2, 2] the new value of x
>>> print(z) # [-1, 0, 1], the old value of x
>>>
>>> ti.atomic_xor(1, x) # will raise TaichiSyntaxError
1076def axes(*x: int): 1077 """Defines a list of axes to be used by a field. 1078 1079 Args: 1080 *x: A list of axes to be activated 1081 1082 Note that Taichi has already provided a set of commonly used axes. For example, 1083 `ti.ij` is just `axes(0, 1)` under the hood. 1084 """ 1085 return [_ti_core.Axis(i) for i in x]
Defines a list of axes to be used by a field.
Args: *x: A list of axes to be activated
Note that Taichi has already provided a set of commonly used axes. For example,
ti.ij is just axes(0, 1) under the hood.
83def bit_cast(obj, dtype): 84 """Copy and cast a scalar to a specified data type with its underlying 85 bits preserved. Must be called in taichi scope. 86 87 This function is equivalent to `reinterpret_cast` in C++. 88 89 Args: 90 obj (:mod:`~taichi.types.primitive_types`): Input scalar. 91 92 dtype (:mod:`~taichi.types.primitive_types`): Target data type, must have \ 93 the same precision bits as the input (hence `f32` -> `f64` is not allowed). 94 95 Returns: 96 A copy of `obj`, casted to the specified data type `dtype`. 97 98 Example:: 99 100 >>> @ti.kernel 101 >>> def test(): 102 >>> x = 3.14 103 >>> y = ti.bit_cast(x, ti.i32) 104 >>> print(y) # 1078523331 105 >>> 106 >>> z = ti.bit_cast(y, ti.f32) 107 >>> print(z) # 3.14 108 """ 109 dtype = cook_dtype(dtype) 110 if is_taichi_class(obj): 111 raise ValueError("Cannot apply bit_cast on Taichi classes") 112 else: 113 return expr.Expr(_ti_core.bits_cast(expr.Expr(obj).ptr, dtype))
Copy and cast a scalar to a specified data type with its underlying bits preserved. Must be called in taichi scope.
This function is equivalent to reinterpret_cast in C++.
Args:
obj (~taichi.types.primitive_types): Input scalar.
dtype (`~taichi.types.primitive_types`): Target data type, must have the same precision bits as the input (hence `f32` -> `f64` is not allowed).
Returns:
A copy of obj, casted to the specified data type dtype.
Example::
>>> @ti.kernel
>>> def test():
>>> x = 3.14
>>> y = ti.bit_cast(x, ti.i32)
>>> print(y) # 1078523331
>>>
>>> z = ti.bit_cast(y, ti.f32)
>>> print(z) # 3.14
1014@taichi_scope 1015def bit_shr(x1, x2): 1016 """Elements in `x1` shifted to the right by number of bits in `x2`. 1017 Both `x1`, `x2` must have integer type. 1018 1019 Args: 1020 x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1021 Input data. 1022 x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1023 Number of bits to remove at the right of `x1`. 1024 1025 Returns: 1026 Return `x1` with bits shifted `x2` times to the right. 1027 This is a scalar if both `x1` and `x2` are scalars. 1028 1029 Example:: 1030 >>> @ti.kernel 1031 >>> def main(): 1032 >>> x = ti.Matrix([7, 8]) 1033 >>> y = ti.Matrix([1, 2]) 1034 >>> print(ti.bit_shr(x, y)) 1035 >>> 1036 >>> main() 1037 [3, 2] 1038 """ 1039 return _binary_operation(_ti_core.expr_bit_shr, _bt_ops_mod.rshift, x1, x2)
Elements in x1 shifted to the right by number of bits in x2.
Both x1, x2 must have integer type.
Args:
x1 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input data.
x2 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Number of bits to remove at the right of x1.
Returns:
Return x1 with bits shifted x2 times to the right.
This is a scalar if both x1 and x2 are scalars.
Example::
@ti.kernel def main(): x = ti.Matrix([7, 8]) y = ti.Matrix([1, 2]) print(ti.bit_shr(x, y))
main() [3, 2]
482def block_local(*args): 483 """Hints Taichi to cache the fields and to enable the BLS optimization. 484 485 Please visit https://docs.taichi-lang.org/docs/performance 486 for how BLS is used. 487 488 Args: 489 *args (List[Field]): A list of sparse Taichi fields. 490 """ 491 if impl.current_cfg().opt_level == 0: 492 _logging.warn("""opt_level = 1 is enforced to enable bls analysis.""") 493 impl.current_cfg().opt_level = 1 494 for a in args: 495 for v in a._get_field_members(): 496 get_runtime().compiling_callable.ast_builder().insert_snode_access_flag( 497 _ti_core.SNodeAccessFlag.block_local, v.ptr 498 )
Hints Taichi to cache the fields and to enable the BLS optimization.
Please visit https://docs.taichi-lang.org/docs/performance for how BLS is used.
Args: *args (List[Field]): A list of sparse Taichi fields.
52def cast(obj, dtype): 53 """Copy and cast a scalar or a matrix to a specified data type. 54 Must be called in Taichi scope. 55 56 Args: 57 obj (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 58 Input scalar or matrix. 59 60 dtype (:mod:`~taichi.types.primitive_types`): A primitive type defined in :mod:`~taichi.types.primitive_types`. 61 62 Returns: 63 A copy of `obj`, casted to the specified data type `dtype`. 64 65 Example:: 66 67 >>> @ti.kernel 68 >>> def test(): 69 >>> x = ti.Matrix([0, 1, 2], ti.i32) 70 >>> y = ti.cast(x, ti.f32) 71 >>> print(y) 72 >>> 73 >>> test() 74 [0.0, 1.0, 2.0] 75 """ 76 dtype = cook_dtype(dtype) 77 if is_taichi_class(obj): 78 # TODO: unify with element_wise_unary 79 return obj.cast(dtype) 80 return expr.Expr(_ti_core.value_cast(expr.Expr(obj).ptr, dtype))
Copy and cast a scalar or a matrix to a specified data type. Must be called in Taichi scope.
Args:
obj (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
dtype (`~taichi.types.primitive_types`): A primitive type defined in `~taichi.types.primitive_types`.
Returns:
A copy of obj, casted to the specified data type dtype.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([0, 1, 2], ti.i32)
>>> y = ti.cast(x, ti.f32)
>>> print(y)
>>>
>>> test()
[0.0, 1.0, 2.0]
373def ceil(x, dtype=None): 374 """Return the ceiling of the input, element-wise. 375 376 The ceil of the scalar `x` is the smallest integer `k`, such that `k >= x`. 377 378 Args: 379 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 380 Input scalar or matrix. 381 382 dtype: (:mod:`~taichi.types.primitive_types`): the returned type, default to `None`. If \ 383 set to `None` the retuned value will have the same type with `x`. 384 385 Returns: 386 The ceiling of each element in `x`, with return value type `dtype`. 387 388 Example:: 389 390 >>> @ti.kernel 391 >>> def test(): 392 >>> x = ti.Matrix([3.14, -1.5]) 393 >>> y = ti.ceil(x) 394 >>> print(y) # [4.0, -1.0] 395 """ 396 result = _ceil(x) 397 if dtype is not None: 398 result = cast(result, dtype) 399 return result
Return the ceiling of the input, element-wise.
The ceil of the scalar x is the smallest integer k, such that k >= x.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
dtype: (`~taichi.types.primitive_types`): the returned type, default to `None`. If set to `None` the retuned value will have the same type with `x`.
Returns:
The ceiling of each element in x, with return value type dtype.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([3.14, -1.5])
>>> y = ti.ceil(x)
>>> print(y) # [4.0, -1.0]
199def cos(x): 200 """Trigonometric cosine, element-wise. 201 202 Args: 203 x (Union[:mod:`~taichi.type.primitive_types`, :class:`~taichi.Matrix`]): \ 204 Angle, in radians. 205 206 Returns: 207 The cosine of each element of `x`. 208 209 Example:: 210 211 >>> from math import pi 212 >>> x = ti.Matrix([-pi, 0, pi/2.]) 213 >>> ti.cos(x) 214 [-1., 1., 0.] 215 """ 216 return _unary_operation(_ti_core.expr_cos, np.cos, x)
Trigonometric cosine, element-wise.
Args:
x (Union[~taichi.type.primitive_types, ~taichi.Matrix]): Angle, in radians.
Returns:
The cosine of each element of x.
Example::
>>> from math import pi
>>> x = ti.Matrix([-pi, 0, pi/2.])
>>> ti.cos(x)
[-1., 1., 0.]
1465def data_oriented(cls): 1466 """Marks a class as Taichi compatible. 1467 1468 To allow for modularized code, Taichi provides this decorator so that 1469 Taichi kernels can be defined inside a class. 1470 1471 See also https://docs.taichi-lang.org/docs/odop 1472 1473 Example:: 1474 1475 >>> @ti.data_oriented 1476 >>> class TiArray: 1477 >>> def __init__(self, n): 1478 >>> self.x = ti.field(ti.f32, shape=n) 1479 >>> 1480 >>> @ti.kernel 1481 >>> def inc(self): 1482 >>> for i in self.x: 1483 >>> self.x[i] += 1.0 1484 >>> 1485 >>> a = TiArray(32) 1486 >>> a.inc() 1487 1488 Args: 1489 cls (Class): the class to be decorated 1490 1491 Returns: 1492 The decorated class. 1493 """ 1494 1495 def _getattr(self, item): 1496 method = cls.__dict__.get(item, None) 1497 is_property = method.__class__ == property 1498 is_staticmethod = method.__class__ == staticmethod 1499 if is_property: 1500 x = method.fget 1501 else: 1502 x = super(cls, self).__getattribute__(item) 1503 if hasattr(x, "_is_wrapped_kernel"): 1504 if inspect.ismethod(x): 1505 wrapped = x.__func__ 1506 else: 1507 wrapped = x 1508 assert isinstance(wrapped, (BoundTaichiCallable, TaichiCallable)) 1509 wrapped._is_staticmethod = is_staticmethod 1510 if wrapped._is_classkernel: 1511 ret = _BoundedDifferentiableMethod(self, wrapped) 1512 ret.__name__ = wrapped.__name__ # type: ignore 1513 if is_property: 1514 return ret() 1515 return ret 1516 if is_property: 1517 return x(self) 1518 return x 1519 1520 cls.__getattribute__ = _getattr 1521 cls._data_oriented = True 1522 1523 return cls
Marks a class as Taichi compatible.
To allow for modularized code, Taichi provides this decorator so that Taichi kernels can be defined inside a class.
See also https://docs.taichi-lang.org/docs/odop
Example::
>>> @ti.data_oriented
>>> class TiArray:
>>> def __init__(self, n):
>>> self.x = ti.field(ti.f32, shape=n)
>>>
>>> @ti.kernel
>>> def inc(self):
>>> for i in self.x:
>>> self.x[i] += 1.0
>>>
>>> a = TiArray(32)
>>> a.inc()
Args: cls (Class): the class to be decorated
Returns: The decorated class.
809def dataclass(cls): 810 """Converts a class with field annotations and methods into a taichi struct type. 811 812 This will return a normal custom struct type, with the functions added to it. 813 Struct fields can be generated in the normal way from the struct type. 814 Functions in the class can be run on the struct instance. 815 816 This class decorator inspects the class for annotations and methods and 817 1. Sets the annotations as fields for the struct 818 2. Attaches the methods to the struct type 819 820 Example:: 821 822 >>> @ti.dataclass 823 >>> class Sphere: 824 >>> center: vec3 825 >>> radius: ti.f32 826 >>> 827 >>> @ti.func 828 >>> def area(self): 829 >>> return 4 * 3.14 * self.radius * self.radius 830 >>> 831 >>> my_spheres = Sphere.field(shape=(n, )) 832 >>> my_sphere[2].area() 833 834 Args: 835 cls (Class): the class with annotations and methods to convert to a struct 836 837 Returns: 838 A taichi struct with the annotations as fields 839 and methods from the class attached. 840 """ 841 # save the annotation fields for the struct 842 fields = getattr(cls, "__annotations__", {}) 843 # raise error if there are default values 844 for k in fields.keys(): 845 if hasattr(cls, k): 846 raise TaichiSyntaxError("Default value in @dataclass is not supported.") 847 # get the class methods to be attached to the struct types 848 fields["__struct_methods"] = { 849 attribute: getattr(cls, attribute) 850 for attribute in dir(cls) 851 if callable(getattr(cls, attribute)) and not attribute.startswith("__") 852 } 853 return StructType(**fields)
Converts a class with field annotations and methods into a taichi struct type.
This will return a normal custom struct type, with the functions added to it. Struct fields can be generated in the normal way from the struct type. Functions in the class can be run on the struct instance.
This class decorator inspects the class for annotations and methods and 1. Sets the annotations as fields for the struct 2. Attaches the methods to the struct type
Example::
>>> @ti.dataclass
>>> class Sphere:
>>> center: vec3
>>> radius: ti.f32
>>>
>>> @ti.func
>>> def area(self):
>>> return 4 * 3.14 * self.radius * self.radius
>>>
>>> my_spheres = Sphere.field(shape=(n, ))
>>> my_sphere[2].area()
Args: cls (Class): the class with annotations and methods to convert to a struct
Returns: A taichi struct with the annotations as fields and methods from the class attached.
426def deactivate(node, indices): 427 """Explicitly deactivate a cell of `node` at location `indices`. 428 429 After deactivation, the Taichi runtime automatically recycles and zero-fills 430 the memory of the deactivated cell. 431 432 Args: 433 node (:class:`~taichi.SNode`): Must be a pointer, hash or bitmasked node. 434 indices (Union[int, :class:`~taichi.Vector`]): the indices to deactivate. 435 """ 436 impl.get_runtime().compiling_callable.ast_builder().insert_deactivate( 437 node._snode.ptr, expr.make_expr_group(indices), _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()) 438 )
Explicitly deactivate a cell of node at location indices.
After deactivation, the Taichi runtime automatically recycles and zero-fills the memory of the deactivated cell.
Args:
node (~taichi.SNode): Must be a pointer, hash or bitmasked node.
indices (Union[int, ~taichi.Vector]): the indices to deactivate.
617def deactivate_all_snodes(): 618 """Recursively deactivate all SNodes.""" 619 for root_fb in FieldsBuilder._finalized_roots(): 620 root_fb.deactivate_all()
Recursively deactivate all SNodes.
457def exp(x): 458 """Compute the exponential of all elements in `x`, element-wise. 459 460 Args: 461 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 462 Input scalar or matrix. 463 464 Returns: 465 Element-wise exponential of `x`. 466 467 Example:: 468 469 >>> @ti.kernel 470 >>> def test(): 471 >>> x = ti.Matrix([-1.0, 0.0, 1.0]) 472 >>> y = ti.exp(x) 473 >>> print(y) 474 >>> 475 >>> test() 476 [0.367879, 1.000000, 2.718282] 477 """ 478 return _unary_operation(_ti_core.expr_exp, np.exp, x)
Compute the exponential of all elements in x, element-wise.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
Returns:
Element-wise exponential of x.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([-1.0, 0.0, 1.0])
>>> y = ti.exp(x)
>>> print(y)
>>>
>>> test()
[0.367879, 1.000000, 2.718282]
805@python_scope 806def field(dtype, *args, **kwargs): 807 """Defines a Taichi field. 808 809 A Taichi field can be viewed as an abstract N-dimensional array, hiding away 810 the complexity of how its underlying :class:`~taichi.lang.snode.SNode` are 811 actually defined. The data in a Taichi field can be directly accessed by 812 a Taichi :func:`~taichi.lang.kernel_impl.kernel`. 813 814 See also https://docs.taichi-lang.org/docs/field 815 816 Args: 817 dtype (DataType): data type of the field. Note it can be vector or matrix types as well. 818 shape (Union[int, tuple[int]], optional): shape of the field. 819 order (str, optional): order of the shape laid out in memory. 820 name (str, optional): name of the field. 821 offset (Union[int, tuple[int]], optional): offset of the field domain. 822 needs_grad (bool, optional): whether this field participates in autodiff (reverse mode) 823 and thus needs an adjoint field to store the gradients. 824 needs_dual (bool, optional): whether this field participates in autodiff (forward mode) 825 and thus needs an dual field to store the gradients. 826 827 Example:: 828 829 The code below shows how a Taichi field can be declared and defined:: 830 831 >>> x1 = ti.field(ti.f32, shape=(16, 8)) 832 >>> # Equivalently 833 >>> x2 = ti.field(ti.f32) 834 >>> ti.root.dense(ti.ij, shape=(16, 8)).place(x2) 835 >>> 836 >>> x3 = ti.field(ti.f32, shape=(16, 8), order='ji') 837 >>> # Equivalently 838 >>> x4 = ti.field(ti.f32) 839 >>> ti.root.dense(ti.j, shape=8).dense(ti.i, shape=16).place(x4) 840 >>> 841 >>> x5 = ti.field(ti.math.vec3, shape=(16, 8)) 842 843 """ 844 if isinstance(dtype, MatrixType): 845 if dtype.ndim == 1: 846 return Vector.field(dtype.n, dtype.dtype, *args, **kwargs) 847 return Matrix.field(dtype.n, dtype.m, dtype.dtype, *args, **kwargs) 848 return _field(dtype, *args, **kwargs)
Defines a Taichi field.
A Taichi field can be viewed as an abstract N-dimensional array, hiding away
the complexity of how its underlying ~taichi.lang.snode.SNode are
actually defined. The data in a Taichi field can be directly accessed by
a Taichi ~taichi.lang.kernel_impl.kernel().
See also https://docs.taichi-lang.org/docs/field
Args: dtype (DataType): data type of the field. Note it can be vector or matrix types as well. shape (Union[int, tuple[int]], optional): shape of the field. order (str, optional): order of the shape laid out in memory. name (str, optional): name of the field. offset (Union[int, tuple[int]], optional): offset of the field domain. needs_grad (bool, optional): whether this field participates in autodiff (reverse mode) and thus needs an adjoint field to store the gradients. needs_dual (bool, optional): whether this field participates in autodiff (forward mode) and thus needs an dual field to store the gradients.
Example::
The code below shows how a Taichi field can be declared and defined::
>>> x1 = ti.field(ti.f32, shape=(16, 8))
>>> # Equivalently
>>> x2 = ti.field(ti.f32)
>>> ti.root.dense(ti.ij, shape=(16, 8)).place(x2)
>>>
>>> x3 = ti.field(ti.f32, shape=(16, 8), order='ji')
>>> # Equivalently
>>> x4 = ti.field(ti.f32)
>>> ti.root.dense(ti.j, shape=8).dense(ti.i, shape=16).place(x4)
>>>
>>> x5 = ti.field(ti.math.vec3, shape=(16, 8))
342def floor(x, dtype=None): 343 """Return the floor of the input, element-wise. 344 The floor of the scalar `x` is the largest integer `k`, such that `k <= x`. 345 346 Args: 347 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 348 Input scalar or matrix. 349 350 dtype: (:mod:`~taichi.types.primitive_types`): the returned type, default to `None`. If \ 351 set to `None` the retuned value will have the same type with `x`. 352 353 Returns: 354 The floor of each element in `x`, with return value type `dtype`. 355 356 Example:: 357 >>> @ti.kernel 358 >>> def test(): 359 >>> x = ti.Matrix([-1.1, 2.2, 3.]) 360 >>> y = ti.floor(x, ti.f64) 361 >>> print(y) # [-2.000000000000, 2.000000000000, 3.000000000000] 362 """ 363 result = _floor(x) 364 if dtype is not None: 365 result = cast(result, dtype) 366 return result
Return the floor of the input, element-wise.
The floor of the scalar x is the largest integer k, such that k <= x.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
dtype: (`~taichi.types.primitive_types`): the returned type, default to `None`. If set to `None` the retuned value will have the same type with `x`.
Returns:
The floor of each element in x, with return value type dtype.
Example::
@ti.kernel def test(): x = ti.Matrix([-1.1, 2.2, 3.]) y = ti.floor(x, ti.f64) print(y) # [-2.000000000000, 2.000000000000, 3.000000000000]
192def func(fn: Callable, is_real_function: bool = False) -> TaichiCallable: 193 """Marks a function as callable in Taichi-scope. 194 195 This decorator transforms a Python function into a Taichi one. Taichi 196 will JIT compile it into native instructions. 197 198 Args: 199 fn (Callable): The Python function to be decorated 200 is_real_function (bool): Whether the function is a real function 201 202 Returns: 203 Callable: The decorated function 204 205 Example:: 206 207 >>> @ti.func 208 >>> def foo(x): 209 >>> return x + 2 210 >>> 211 >>> @ti.kernel 212 >>> def run(): 213 >>> print(foo(40)) # 42 214 """ 215 is_classfunc = _inside_class(level_of_class_stackframe=3 + is_real_function) 216 217 fun = Func(fn, _classfunc=is_classfunc, is_real_function=is_real_function) 218 taichi_callable = TaichiCallable(fn, fun) 219 taichi_callable._is_taichi_function = True 220 taichi_callable._is_real_function = is_real_function 221 return taichi_callable
Marks a function as callable in Taichi-scope.
This decorator transforms a Python function into a Taichi one. Taichi will JIT compile it into native instructions.
Args: fn (Callable): The Python function to be decorated is_real_function (bool): Whether the function is a real function
Returns: Callable: The decorated function
Example::
>>> @ti.func
>>> def foo(x):
>>> return x + 2
>>>
>>> @ti.kernel
>>> def run():
>>> print(foo(40)) # 42
459def get_addr(f, indices): 460 """Query the memory address (on CUDA/x64) of field `f` at index `indices`. 461 462 Currently, this function can only be called inside a taichi kernel. 463 464 Args: 465 f (Union[:class:`~taichi.Field`, :class:`~taichi.MatrixField`]): Input taichi field for memory address query. 466 indices (Union[int, :class:`~taichi.Vector`]): The specified field indices of the query. 467 468 Returns: 469 ti.u64: The memory address of `f[indices]`. 470 """ 471 return expr.Expr( 472 impl.get_runtime() 473 .compiling_callable.ast_builder() 474 .expr_snode_get_addr(f._snode.ptr, expr.make_expr_group(indices)), 475 dbg_info=_ti_core.DebugInfo(impl.get_runtime().get_current_src_info()), 476 )
Query the memory address (on CUDA/x64) of field f at index indices.
Currently, this function can only be called inside a taichi kernel.
Args:
f (Union[~taichi.Field, ~taichi.MatrixField]): Input taichi field for memory address query.
indices (Union[int, ~taichi.Vector]): The specified field indices of the query.
Returns:
ti.u64: The memory address of f[indices].
683def global_thread_idx(): 684 """Returns the global thread id of this running thread, 685 only available for cpu and cuda backends. 686 687 For cpu backends this is equal to the cpu thread id, 688 For cuda backends this is equal to `block_id * block_dim + thread_id`. 689 690 Example:: 691 692 >>> f = ti.field(ti.f32, shape=(16, 16)) 693 >>> @ti.kernel 694 >>> def test(): 695 >>> for i in ti.grouped(f): 696 >>> print(ti.global_thread_idx()) 697 >>> 698 test() 699 """ 700 return impl.get_runtime().compiling_callable.ast_builder().insert_thread_idx_expr()
Returns the global thread id of this running thread, only available for cpu and cuda backends.
For cpu backends this is equal to the cpu thread id,
For cuda backends this is equal to block_id * block_dim + thread_id.
Example::
>>> f = ti.field(ti.f32, shape=(16, 16))
>>> @ti.kernel
>>> def test():
>>> for i in ti.grouped(f):
>>> print(ti.global_thread_idx())
>>>
test()
1169@taichi_scope 1170def grouped(x): 1171 """Groups the indices in the iterator returned by `ndrange()` into a 1-D vector. 1172 1173 This is often used when you want to iterate over all indices returned by `ndrange()` 1174 in one `for` loop and a single index. 1175 1176 Args: 1177 x (:func:`~taichi.ndrange`): an iterator object returned by `ti.ndrange`. 1178 1179 Example:: 1180 >>> # without ti.grouped 1181 >>> for I in ti.ndrange(2, 3): 1182 >>> print(I) 1183 prints 0, 1, 2, 3, 4, 5 1184 1185 >>> # with ti.grouped 1186 >>> for I in ti.grouped(ti.ndrange(2, 3)): 1187 >>> print(I) 1188 prints [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2] 1189 """ 1190 if isinstance(x, _Ndrange): 1191 return x.grouped() 1192 return x
Groups the indices in the iterator returned by ndrange() into a 1-D vector.
This is often used when you want to iterate over all indices returned by ndrange()
in one for loop and a single index.
Args:
x (~taichi.ndrange()): an iterator object returned by ti.ndrange.
Example::
without ti.grouped
for I in ti.ndrange(2, 3): print(I) prints 0, 1, 2, 3, 4, 5
>>> # with ti.grouped >>> for I in ti.grouped(ti.ndrange(2, 3)): >>> print(I) prints [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]
325def init( 326 arch=None, 327 default_fp=None, 328 default_ip=None, 329 _test_mode=False, 330 enable_fallback=True, 331 require_version=None, 332 **kwargs, 333): 334 """Initializes the Taichi runtime. 335 336 This should always be the entry point of your Taichi program. Most 337 importantly, it sets the backend used throughout the program. 338 339 Args: 340 arch: Backend to use. This is usually :const:`~taichi.lang.cpu` or :const:`~taichi.lang.gpu`. 341 default_fp (Optional[type]): Default floating-point type. 342 default_ip (Optional[type]): Default integral type. 343 require_version (Optional[string]): A version string. 344 **kwargs: Taichi provides highly customizable compilation through 345 ``kwargs``, which allows for fine grained control of Taichi compiler 346 behavior. Below we list some of the most frequently used ones. For a 347 complete list, please check out 348 https://github.com/taichi-dev/taichi/blob/master/taichi/program/compile_config.h. 349 350 * ``cpu_max_num_threads`` (int): Sets the number of threads used by the CPU thread pool. 351 * ``debug`` (bool): Enables the debug mode, under which Taichi does a few more things like boundary checks. 352 * ``print_ir`` (bool): Prints the CHI IR of the Taichi kernels. 353 *``offline_cache`` (bool): Enables offline cache of the compiled kernels. Default to True. When this is enabled Taichi will cache compiled kernel on your local disk to accelerate future calls. 354 *``random_seed`` (int): Sets the seed of the random generator. The default is 0. 355 """ 356 # Check version for users every 7 days if not disabled by users. 357 _version_check.start_version_check_thread() 358 359 # FIXME(https://github.com/taichi-dev/taichi/issues/4811): save the current working directory since it may be 360 # changed by the Vulkan backend initialization on OS X. 361 current_dir = os.getcwd() 362 363 # Check if installed version meets the requirements. 364 if require_version is not None: 365 check_require_version(require_version) 366 367 if "default_up" in kwargs: 368 raise KeyError("'default_up' is always the unsigned type of 'default_ip'. Please set 'default_ip' instead.") 369 # Make a deepcopy in case these args reference to items from ti.cfg, which are 370 # actually references. If no copy is made and the args are indeed references, 371 # ti.reset() could override the args to their default values. 372 default_fp = _deepcopy(default_fp) 373 default_ip = _deepcopy(default_ip) 374 kwargs = _deepcopy(kwargs) 375 reset() 376 377 cfg = impl.default_cfg() 378 cfg.offline_cache = True # Enable offline cache in frontend instead of C++ side 379 380 spec_cfg = _SpecialConfig() 381 env_comp = _EnvironmentConfigurator(kwargs, cfg) 382 env_spec = _EnvironmentConfigurator(kwargs, spec_cfg) 383 384 # configure default_fp/ip: 385 # TODO: move these stuff to _SpecialConfig too: 386 env_default_fp = os.environ.get("TI_DEFAULT_FP") 387 if env_default_fp: 388 if default_fp is not None: 389 _ti_core.warn( 390 f'Environment variable TI_DEFAULT_FP={env_default_fp} overridden by ti.init argument "default_fp"' 391 ) 392 elif env_default_fp == "32": 393 default_fp = f32 394 elif env_default_fp == "64": 395 default_fp = f64 396 elif env_default_fp is not None: 397 raise ValueError(f"Invalid TI_DEFAULT_FP={env_default_fp}, should be 32 or 64") 398 399 env_default_ip = os.environ.get("TI_DEFAULT_IP") 400 if env_default_ip: 401 if default_ip is not None: 402 _ti_core.warn( 403 f'Environment variable TI_DEFAULT_IP={env_default_ip} overridden by ti.init argument "default_ip"' 404 ) 405 elif env_default_ip == "32": 406 default_ip = i32 407 elif env_default_ip == "64": 408 default_ip = i64 409 elif env_default_ip is not None: 410 raise ValueError(f"Invalid TI_DEFAULT_IP={env_default_ip}, should be 32 or 64") 411 412 if default_fp is not None: 413 impl.get_runtime().set_default_fp(default_fp) 414 if default_ip is not None: 415 impl.get_runtime().set_default_ip(default_ip) 416 417 # submodule configurations (spec_cfg): 418 env_spec.add("log_level", str) 419 env_spec.add("gdb_trigger") 420 env_spec.add("short_circuit_operators") 421 env_spec.add("print_full_traceback") 422 env_spec.add("unrolling_limit") 423 424 # compiler configurations (ti.cfg): 425 for key in dir(cfg): 426 if key in ["arch", "default_fp", "default_ip"]: 427 continue 428 _cast = type(getattr(cfg, key)) 429 if _cast is bool: 430 _cast = None 431 env_comp.add(key, _cast) 432 433 unexpected_keys = kwargs.keys() 434 435 if len(unexpected_keys): 436 raise KeyError(f'Unrecognized keyword argument(s) for ti.init: {", ".join(unexpected_keys)}') 437 438 # dispatch configurations that are not in ti.cfg: 439 if not _test_mode: 440 _ti_core.set_core_trigger_gdb_when_crash(spec_cfg.gdb_trigger) 441 impl.get_runtime().short_circuit_operators = spec_cfg.short_circuit_operators 442 impl.get_runtime().print_full_traceback = spec_cfg.print_full_traceback 443 impl.get_runtime().unrolling_limit = spec_cfg.unrolling_limit 444 _logging.set_logging_level(spec_cfg.log_level.lower()) 445 446 # select arch (backend): 447 env_arch = os.environ.get("TI_ARCH") 448 if env_arch is not None: 449 _logging.info(f"Following TI_ARCH setting up for arch={env_arch}") 450 arch = _ti_core.arch_from_name(env_arch) 451 cfg.arch = adaptive_arch_select(arch, enable_fallback) 452 print(f"[Taichi] Starting on arch={_ti_core.arch_name(cfg.arch)}") 453 454 if _test_mode: 455 return spec_cfg 456 457 get_default_kernel_profiler().set_kernel_profiler_mode(cfg.kernel_profiler) 458 459 # create a new program: 460 impl.get_runtime().create_program() 461 462 _logging.trace("Materializing runtime...") 463 impl.get_runtime().prog.materialize_runtime() 464 465 impl._root_fb = _snode.FieldsBuilder() 466 467 if cfg.debug: 468 impl.get_runtime()._register_signal_handlers() 469 470 # Recover the current working directory (https://github.com/taichi-dev/taichi/issues/4811) 471 os.chdir(current_dir) 472 return None
Initializes the Taichi runtime.
This should always be the entry point of your Taichi program. Most importantly, it sets the backend used throughout the program.
Args:
arch: Backend to use. This is usually ~taichi.lang.cpu or ~taichi.lang.gpu.
default_fp (Optional[type]): Default floating-point type.
default_ip (Optional[type]): Default integral type.
require_version (Optional[string]): A version string.
**kwargs: Taichi provides highly customizable compilation through
kwargs, which allows for fine grained control of Taichi compiler
behavior. Below we list some of the most frequently used ones. For a
complete list, please check out
https://github.com/taichi-dev/taichi/blob/master/taichi/program/compile_config.h.
* ``cpu_max_num_threads`` (int): Sets the number of threads used by the CPU thread pool.
* ``debug`` (bool): Enables the debug mode, under which Taichi does a few more things like boundary checks.
* ``print_ir`` (bool): Prints the CHI IR of the Taichi kernels.
*``offline_cache`` (bool): Enables offline cache of the compiled kernels. Default to True. When this is enabled Taichi will cache compiled kernel on your local disk to accelerate future calls.
*``random_seed`` (int): Sets the seed of the random generator. The default is 0.
395def is_active(node, indices): 396 """Explicitly query whether a cell in a SNode `node` at location 397 `indices` is active or not. 398 399 Args: 400 node (:class:`~taichi.SNode`): Must be a pointer, hash or bitmasked node. 401 indices (Union[int, list, :class:`~taichi.Vector`]): the indices to visit. 402 403 Returns: 404 bool: the cell `node[indices]` is active or not. 405 """ 406 return expr.Expr( 407 impl.get_runtime() 408 .compiling_callable.ast_builder() 409 .expr_snode_is_active(node._snode.ptr, expr.make_expr_group(indices)), 410 dbg_info=_ti_core.DebugInfo(impl.get_runtime().get_current_src_info()), 411 )
Explicitly query whether a cell in a SNode node at location
indices is active or not.
Args:
node (~taichi.SNode): Must be a pointer, hash or bitmasked node.
indices (Union[int, list, ~taichi.Vector]): the indices to visit.
Returns:
bool: the cell node[indices] is active or not.
1406def kernel(fn: Callable): 1407 """Marks a function as a Taichi kernel. 1408 1409 A Taichi kernel is a function written in Python, and gets JIT compiled by 1410 Taichi into native CPU/GPU instructions (e.g. a series of CUDA kernels). 1411 The top-level ``for`` loops are automatically parallelized, and distributed 1412 to either a CPU thread pool or massively parallel GPUs. 1413 1414 Kernel's gradient kernel would be generated automatically by the AutoDiff system. 1415 1416 See also https://docs.taichi-lang.org/docs/syntax#kernel. 1417 1418 Args: 1419 fn (Callable): the Python function to be decorated 1420 1421 Returns: 1422 Callable: The decorated function 1423 1424 Example:: 1425 1426 >>> x = ti.field(ti.i32, shape=(4, 8)) 1427 >>> 1428 >>> @ti.kernel 1429 >>> def run(): 1430 >>> # Assigns all the elements of `x` in parallel. 1431 >>> for i in x: 1432 >>> x[i] = i 1433 """ 1434 return _kernel_impl(fn, level_of_class_stackframe=3)
Marks a function as a Taichi kernel.
A Taichi kernel is a function written in Python, and gets JIT compiled by
Taichi into native CPU/GPU instructions (e.g. a series of CUDA kernels).
The top-level for loops are automatically parallelized, and distributed
to either a CPU thread pool or massively parallel GPUs.
Kernel's gradient kernel would be generated automatically by the AutoDiff system.
See also https://docs.taichi-lang.org/docs/syntax#kernel.
Args: fn (Callable): the Python function to be decorated
Returns: Callable: The decorated function
Example::
>>> x = ti.field(ti.i32, shape=(4, 8))
>>>
>>> @ti.kernel
>>> def run():
>>> # Assigns all the elements of `x` in parallel.
>>> for i in x:
>>> x[i] = i
441def length(node, indices): 442 """Return the length of the dynamic SNode `node` at index `indices`. 443 444 Args: 445 node (:class:`~taichi.SNode`): a dynamic SNode. 446 indices (Union[int, :class:`~taichi.Vector`]): the indices to query. 447 448 Returns: 449 int: the length of cell `node[indices]`. 450 """ 451 return expr.Expr( 452 impl.get_runtime() 453 .compiling_callable.ast_builder() 454 .expr_snode_length(node._snode.ptr, expr.make_expr_group(indices)), 455 dbg_info=_ti_core.DebugInfo(impl.get_runtime().get_current_src_info()), 456 )
Return the length of the dynamic SNode node at index indices.
Args:
node (~taichi.SNode): a dynamic SNode.
indices (Union[int, ~taichi.Vector]): the indices to query.
Returns:
int: the length of cell node[indices].
481def log(x): 482 """Compute the natural logarithm, element-wise. 483 484 The natural logarithm `log` is the inverse of the exponential function, 485 so that `log(exp(x)) = x`. The natural logarithm is logarithm in base `e`. 486 487 Args: 488 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 489 Input scalar or matrix. 490 491 Returns: 492 The natural logarithm of `x`, element-wise. 493 494 Example:: 495 496 >>> @ti.kernel 497 >>> def test(): 498 >>> x = ti.Vector([-1.0, 0.0, 1.0]) 499 >>> y = ti.log(x) 500 >>> print(y) 501 >>> 502 >>> test() 503 [-nan, -inf, 0.000000] 504 """ 505 return _unary_operation(_ti_core.expr_log, np.log, x)
Compute the natural logarithm, element-wise.
The natural logarithm log is the inverse of the exponential function,
so that log(exp(x)) = x. The natural logarithm is logarithm in base e.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
Returns:
The natural logarithm of x, element-wise.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([-1.0, 0.0, 1.0])
>>> y = ti.log(x)
>>> print(y)
>>>
>>> test()
[-nan, -inf, 0.000000]
614def loop_config( 615 *, 616 block_dim=None, 617 serialize=False, 618 parallelize=None, 619 block_dim_adaptive=True, 620 bit_vectorize=False, 621): 622 """Sets directives for the next loop 623 624 Args: 625 block_dim (int): The number of threads in a block on GPU 626 serialize (bool): Whether to let the for loop execute serially, `serialize=True` equals to `parallelize=1` 627 parallelize (int): The number of threads to use on CPU 628 block_dim_adaptive (bool): Whether to allow backends set block_dim adaptively, enabled by default 629 bit_vectorize (bool): Whether to enable bit vectorization of struct fors on quant_arrays. 630 631 Examples:: 632 633 @ti.kernel 634 def break_in_serial_for() -> ti.i32: 635 a = 0 636 ti.loop_config(serialize=True) 637 for i in range(100): # This loop runs serially 638 a += i 639 if i == 10: 640 break 641 return a 642 643 break_in_serial_for() # returns 55 644 645 n = 128 646 val = ti.field(ti.i32, shape=n) 647 @ti.kernel 648 def fill(): 649 ti.loop_config(parallelize=8, block_dim=16) 650 # If the kernel is run on the CPU backend, 8 threads will be used to run it 651 # If the kernel is run on the CUDA backend, each block will have 16 threads. 652 for i in range(n): 653 val[i] = i 654 655 u1 = ti.types.quant.int(bits=1, signed=False) 656 x = ti.field(dtype=u1) 657 y = ti.field(dtype=u1) 658 cell = ti.root.dense(ti.ij, (128, 4)) 659 cell.quant_array(ti.j, 32).place(x) 660 cell.quant_array(ti.j, 32).place(y) 661 @ti.kernel 662 def copy(): 663 ti.loop_config(bit_vectorize=True) 664 # 32 bits, instead of 1 bit, will be copied at a time 665 for i, j in x: 666 y[i, j] = x[i, j] 667 """ 668 if block_dim is not None: 669 _block_dim(block_dim) 670 671 if serialize: 672 _parallelize(1) 673 elif parallelize is not None: 674 _parallelize(parallelize) 675 676 if not block_dim_adaptive: 677 _block_dim_adaptive(block_dim_adaptive) 678 679 if bit_vectorize: 680 _bit_vectorize()
Sets directives for the next loop
Args:
block_dim (int): The number of threads in a block on GPU
serialize (bool): Whether to let the for loop execute serially, serialize=True equals to parallelize=1
parallelize (int): The number of threads to use on CPU
block_dim_adaptive (bool): Whether to allow backends set block_dim adaptively, enabled by default
bit_vectorize (bool): Whether to enable bit vectorization of struct fors on quant_arrays.
Examples::
@ti.kernel
def break_in_serial_for() -> ti.i32:
a = 0
ti.loop_config(serialize=True)
for i in range(100): # This loop runs serially
a += i
if i == 10:
break
return a
break_in_serial_for() # returns 55
n = 128
val = ti.field(ti.i32, shape=n)
@ti.kernel
def fill():
ti.loop_config(parallelize=8, block_dim=16)
# If the kernel is run on the CPU backend, 8 threads will be used to run it
# If the kernel is run on the CUDA backend, each block will have 16 threads.
for i in range(n):
val[i] = i
u1 = ti.types.quant.int(bits=1, signed=False)
x = ti.field(dtype=u1)
y = ti.field(dtype=u1)
cell = ti.root.dense(ti.ij, (128, 4))
cell.quant_array(ti.j, 32).place(x)
cell.quant_array(ti.j, 32).place(y)
@ti.kernel
def copy():
ti.loop_config(bit_vectorize=True)
# 32 bits, instead of 1 bit, will be copied at a time
for i, j in x:
y[i, j] = x[i, j]
1391def max(*args): # pylint: disable=W0622 1392 """Compute the maximum of the arguments, element-wise. 1393 1394 This function takes no effect on a single argument, even it's array-like. 1395 When there are both scalar and matrix arguments in `args`, the matrices 1396 must have the same shape, and scalars will be broadcasted to the same shape as the matrix. 1397 1398 Args: 1399 args: (List[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1400 The input. 1401 1402 Returns: 1403 Maximum of the inputs. 1404 1405 Example:: 1406 1407 >>> @ti.kernel 1408 >>> def foo(): 1409 >>> x = ti.Vector([0, 1, 2]) 1410 >>> y = ti.Vector([3, 4, 5]) 1411 >>> z = ti.max(x, y, 4) 1412 >>> print(z) # [4, 4, 5] 1413 """ 1414 num_args = len(args) 1415 assert num_args >= 1 1416 if num_args == 1: 1417 return args[0] 1418 if num_args == 2: 1419 return max_impl(args[0], args[1]) 1420 return max_impl(args[0], max(*args[1:]))
Compute the maximum of the arguments, element-wise.
This function takes no effect on a single argument, even it's array-like.
When there are both scalar and matrix arguments in args, the matrices
must have the same shape, and scalars will be broadcasted to the same shape as the matrix.
Args:
args: (List[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns: Maximum of the inputs.
Example::
>>> @ti.kernel
>>> def foo():
>>> x = ti.Vector([0, 1, 2])
>>> y = ti.Vector([3, 4, 5])
>>> z = ti.max(x, y, 4)
>>> print(z) # [4, 4, 5]
501def mesh_local(*args): 502 """Hints the compiler to cache the mesh attributes 503 and to enable the mesh BLS optimization, 504 only available for backends supporting `ti.extension.mesh` and to use with mesh-for loop. 505 506 Related to https://github.com/taichi-dev/taichi/issues/3608 507 508 Args: 509 *args (List[Attribute]): A list of mesh attributes or fields accessed as attributes. 510 511 Examples:: 512 513 # instantiate model 514 mesh_builder = ti.Mesh.tri() 515 mesh_builder.verts.place({ 516 'x' : ti.f32, 517 'y' : ti.f32 518 }) 519 model = mesh_builder.build(meta) 520 521 @ti.kernel 522 def foo(): 523 # hint the compiler to cache mesh vertex attribute `x` and `y`. 524 ti.mesh_local(model.verts.x, model.verts.y) 525 for v0 in model.verts: # mesh-for loop 526 for v1 in v0.verts: 527 v0.x += v1.y 528 """ 529 for a in args: 530 for v in a._get_field_members(): 531 get_runtime().compiling_callable.ast_builder().insert_snode_access_flag( 532 _ti_core.SNodeAccessFlag.mesh_local, v.ptr 533 )
Hints the compiler to cache the mesh attributes
and to enable the mesh BLS optimization,
only available for backends supporting ti.extension.mesh and to use with mesh-for loop.
Related to https://github.com/taichi-dev/taichi/issues/3608
Args: *args (List[Attribute]): A list of mesh attributes or fields accessed as attributes.
Examples::
# instantiate model
mesh_builder = ti.Mesh.tri()
mesh_builder.verts.place({
'x' : ti.f32,
'y' : ti.f32
})
model = mesh_builder.build(meta)
@ti.kernel
def foo():
# hint the compiler to cache mesh vertex attribute `x` and `y`.
ti.mesh_local(model.verts.x, model.verts.y)
for v0 in model.verts: # mesh-for loop
for v1 in v0.verts:
v0.x += v1.y
703def mesh_patch_idx(): 704 """Returns the internal mesh patch id of this running thread, 705 only available for backends supporting `ti.extension.mesh` and to use within mesh-for loop. 706 707 Related to https://github.com/taichi-dev/taichi/issues/3608 708 """ 709 return ( 710 impl.get_runtime() 711 .compiling_callable.ast_builder() 712 .insert_patch_idx_expr(_ti_core.DebugInfo(impl.get_runtime().get_current_src_info())) 713 )
Returns the internal mesh patch id of this running thread,
only available for backends supporting ti.extension.mesh and to use within mesh-for loop.
1423def min(*args): # pylint: disable=W0622 1424 """Compute the minimum of the arguments, element-wise. 1425 1426 This function takes no effect on a single argument, even it's array-like. 1427 When there are both scalar and matrix arguments in `args`, the matrices 1428 must have the same shape, and scalars will be broadcasted to the same shape as the matrix. 1429 1430 Args: 1431 args: (List[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1432 The input. 1433 1434 Returns: 1435 Minimum of the inputs. 1436 1437 Example:: 1438 1439 >>> @ti.kernel 1440 >>> def foo(): 1441 >>> x = ti.Vector([0, 1, 2]) 1442 >>> y = ti.Vector([3, 4, 5]) 1443 >>> z = ti.min(x, y, 1) 1444 >>> print(z) # [0, 1, 1] 1445 """ 1446 num_args = len(args) 1447 assert num_args >= 1 1448 if num_args == 1: 1449 return args[0] 1450 if num_args == 2: 1451 return min_impl(args[0], args[1]) 1452 return min_impl(args[0], min(*args[1:]))
Compute the minimum of the arguments, element-wise.
This function takes no effect on a single argument, even it's array-like.
When there are both scalar and matrix arguments in args, the matrices
must have the same shape, and scalars will be broadcasted to the same shape as the matrix.
Args:
args: (List[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns: Minimum of the inputs.
Example::
>>> @ti.kernel
>>> def foo():
>>> x = ti.Vector([0, 1, 2])
>>> y = ti.Vector([3, 4, 5])
>>> z = ti.min(x, y, 1)
>>> print(z) # [0, 1, 1]
851@python_scope 852def ndarray(dtype, shape, needs_grad=False): 853 """Defines a Taichi ndarray with scalar elements. 854 855 Args: 856 dtype (Union[DataType, MatrixType]): Data type of each element. This can be either a scalar type like ti.f32 or a compound type like ti.types.vector(3, ti.i32). 857 shape (Union[int, tuple[int]]): Shape of the ndarray. 858 859 Example: 860 The code below shows how a Taichi ndarray with scalar elements can be declared and defined:: 861 862 >>> x = ti.ndarray(ti.f32, shape=(16, 8)) # ndarray of shape (16, 8), each element is ti.f32 scalar. 863 >>> vec3 = ti.types.vector(3, ti.i32) 864 >>> y = ti.ndarray(vec3, shape=(10, 2)) # ndarray of shape (10, 2), each element is a vector of 3 ti.i32 scalars. 865 >>> matrix_ty = ti.types.matrix(3, 4, float) 866 >>> z = ti.ndarray(matrix_ty, shape=(4, 5)) # ndarray of shape (4, 5), each element is a matrix of (3, 4) ti.float scalars. 867 """ 868 # primal 869 if isinstance(shape, numbers.Number): 870 shape = (shape,) 871 if not all((isinstance(x, int) or isinstance(x, np.integer)) and x > 0 and x <= 2**31 - 1 for x in shape): 872 raise TaichiRuntimeError(f"{shape} is not a valid shape for ndarray") 873 if dtype in all_types: 874 dt = cook_dtype(dtype) 875 x = ScalarNdarray(dt, shape) 876 elif isinstance(dtype, MatrixType): 877 if dtype.ndim == 1: 878 x = VectorNdarray(dtype.n, dtype.dtype, shape) 879 else: 880 x = MatrixNdarray(dtype.n, dtype.m, dtype.dtype, shape) 881 dt = dtype.dtype 882 else: 883 raise TaichiRuntimeError(f"{dtype} is not supported as ndarray element type") 884 if needs_grad: 885 assert isinstance(dt, DataType) 886 if not _ti_core.is_real(dt): 887 raise TaichiRuntimeError(f"{dt} is not supported for ndarray with `needs_grad=True` or `needs_dual=True`.") 888 x_grad = ndarray(dtype, shape, needs_grad=False) 889 x._set_grad(x_grad) 890 return x
Defines a Taichi ndarray with scalar elements.
Args: dtype (Union[DataType, MatrixType]): Data type of each element. This can be either a scalar type like ti.f32 or a compound type like ti.types.vector(3, ti.i32). shape (Union[int, tuple[int]]): Shape of the ndarray.
Example: The code below shows how a Taichi ndarray with scalar elements can be declared and defined::
>>> x = ti.ndarray(ti.f32, shape=(16, 8)) # ndarray of shape (16, 8), each element is ti.f32 scalar.
>>> vec3 = ti.types.vector(3, ti.i32)
>>> y = ti.ndarray(vec3, shape=(10, 2)) # ndarray of shape (10, 2), each element is a vector of 3 ti.i32 scalars.
>>> matrix_ty = ti.types.matrix(3, 4, float)
>>> z = ti.ndarray(matrix_ty, shape=(4, 5)) # ndarray of shape (4, 5), each element is a matrix of (3, 4) ti.float scalars.
61def ndrange(*args) -> Iterable: 62 """Return an immutable iterator object for looping over multi-dimensional indices. 63 64 This returned set of multi-dimensional indices is the direct product (in the set-theory sense) 65 of n groups of integers, where n equals the number of arguments in the input list, and looks like 66 67 range(x1, y1) x range(x2, y2) x ... x range(xn, yn) 68 69 The k-th argument corresponds to the k-th `range()` factor in the above product, and each 70 argument must be an integer or a pair of two integers. An integer argument n will be interpreted 71 as `range(0, n)`, and a pair of two integers (start, end) will be interpreted as `range(start, end)`. 72 73 You can loop over these multi-dimensonal indices in different ways, see the examples below. 74 75 Args: 76 entries: (int, tuple): Must be either an integer, or a tuple/list of two integers. 77 78 Returns: 79 An immutable iterator object. 80 81 Example:: 82 83 You can loop over 1-D integers in range [start, end), as in native Python 84 85 >>> @ti.kernel 86 >>> def loop_1d(): 87 >>> start = 2 88 >>> end = 5 89 >>> for i in ti.ndrange((start, end)): 90 >>> print(i) # will print 2 3 4 91 92 Note the braces around `(start, end)` in the above code. If without them, 93 the parameter `2` will be interpreted as `range(0, 2)`, `5` will be 94 interpreted as `range(0, 5)`, and you will get a set of 2-D indices which 95 contains 2x5=10 elements, and need two indices i, j to loop over them: 96 97 >>> @ti.kernel 98 >>> def loop_2d(): 99 >>> for i, j in ti.ndrange(2, 5): 100 >>> print(i, j) 101 0 0 102 ... 103 0 4 104 ... 105 1 4 106 107 But you do can use a single index i to loop over these 2-D indices, in this case 108 the indices are returned as a 1-D array `(0, 1, ..., 9)`: 109 110 >>> @ti.kernel 111 >>> def loop_2d_as_1d(): 112 >>> for i in ti.ndrange(2, 5): 113 >>> print(i) 114 will print 0 1 2 3 4 5 6 7 8 9 115 116 In general, you can use any `1 <= k <= n` iterators to loop over a set of n-D 117 indices. For `k=n` all the indices are n-dimensional, and they are returned in 118 lexical order, but for `k<n` iterators the last n-k+1 dimensions will be collapsed into 119 a 1-D array of consecutive integers `(0, 1, 2, ...)` whose length equals the 120 total number of indices in the last n-k+1 dimensions: 121 122 >>> @ti.kernel 123 >>> def loop_3d_as_2d(): 124 >>> # use two iterators to loop over a set of 3-D indices 125 >>> # the last two dimensions for 4, 5 will collapse into 126 >>> # the array [0, 1, 2, ..., 19] 127 >>> for i, j in ti.ndrange(3, 4, 5): 128 >>> print(i, j) 129 will print 0 0, 0 1, ..., 0 19, ..., 2 19. 130 131 A typical usage of `ndrange` is when you want to loop over a tensor and process 132 its entries in parallel. You should avoid writing nested `for` loops here since 133 only top level `for` loops are paralleled in taichi, instead you can use `ndrange` 134 to hold all entries in one top level loop: 135 136 >>> @ti.kernel 137 >>> def loop_tensor(): 138 >>> for row, col, channel in ti.ndrange(image_height, image_width, channels): 139 >>> image[row, col, channel] = ... 140 """ 141 return _Ndrange(*args)
Return an immutable iterator object for looping over multi-dimensional indices.
This returned set of multi-dimensional indices is the direct product (in the set-theory sense) of n groups of integers, where n equals the number of arguments in the input list, and looks like
range(x1, y1) x range(x2, y2) x ... x range(xn, yn)
The k-th argument corresponds to the k-th range() factor in the above product, and each
argument must be an integer or a pair of two integers. An integer argument n will be interpreted
as range(0, n), and a pair of two integers (start, end) will be interpreted as range(start, end).
You can loop over these multi-dimensonal indices in different ways, see the examples below.
Args: entries: (int, tuple): Must be either an integer, or a tuple/list of two integers.
Returns: An immutable iterator object.
Example::
You can loop over 1-D integers in range [start, end), as in native Python
>>> @ti.kernel
>>> def loop_1d():
>>> start = 2
>>> end = 5
>>> for i in ti.ndrange((start, end)):
>>> print(i) # will print 2 3 4
Note the braces around `(start, end)` in the above code. If without them,
the parameter `2` will be interpreted as `range(0, 2)`, `5` will be
interpreted as `range(0, 5)`, and you will get a set of 2-D indices which
contains 2x5=10 elements, and need two indices i, j to loop over them:
>>> @ti.kernel
>>> def loop_2d():
>>> for i, j in ti.ndrange(2, 5):
>>> print(i, j)
0 0
...
0 4
...
1 4
But you do can use a single index i to loop over these 2-D indices, in this case
the indices are returned as a 1-D array `(0, 1, ..., 9)`:
>>> @ti.kernel
>>> def loop_2d_as_1d():
>>> for i in ti.ndrange(2, 5):
>>> print(i)
will print 0 1 2 3 4 5 6 7 8 9
In general, you can use any `1 <= k <= n` iterators to loop over a set of n-D
indices. For `k=n` all the indices are n-dimensional, and they are returned in
lexical order, but for `k<n` iterators the last n-k+1 dimensions will be collapsed into
a 1-D array of consecutive integers `(0, 1, 2, ...)` whose length equals the
total number of indices in the last n-k+1 dimensions:
>>> @ti.kernel
>>> def loop_3d_as_2d():
>>> # use two iterators to loop over a set of 3-D indices
>>> # the last two dimensions for 4, 5 will collapse into
>>> # the array [0, 1, 2, ..., 19]
>>> for i, j in ti.ndrange(3, 4, 5):
>>> print(i, j)
will print 0 0, 0 1, ..., 0 19, ..., 2 19.
A typical usage of `ndrange` is when you want to loop over a tensor and process
its entries in parallel. You should avoid writing nested `for` loops here since
only top level `for` loops are paralleled in taichi, instead you can use `ndrange`
to hold all entries in one top level loop:
>>> @ti.kernel
>>> def loop_tensor():
>>> for row, col, channel in ti.ndrange(image_height, image_width, channels):
>>> image[row, col, channel] = ...
475def no_activate(*args): 476 """Deactivates a SNode pointer.""" 477 assert isinstance(get_runtime().compiling_callable, _ti_core.Kernel) 478 for v in args: 479 get_runtime().compiling_callable.no_activate(v._snode.ptr)
Deactivates a SNode pointer.
1053@taichi_scope 1054def one(x): 1055 """Returns an array of ones with the same shape and type as the input. It's also a scalar 1056 if the input is a scalar. 1057 1058 Args: 1059 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): The input. 1060 1061 Returns: 1062 A new copy of the input but filled with ones. 1063 1064 Example:: 1065 1066 >>> x = ti.Vector([0, 0]) 1067 >>> @ti.kernel 1068 >>> def test(): 1069 >>> y = ti.one(x) 1070 >>> print(y) 1071 [1, 1] 1072 """ 1073 return zero(x) + 1
Returns an array of ones with the same shape and type as the input. It's also a scalar if the input is a scalar.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns: A new copy of the input but filled with ones.
Example::
>>> x = ti.Vector([0, 0])
>>> @ti.kernel
>>> def test():
>>> y = ti.one(x)
>>> print(y)
[1, 1]
683def pow(base, exponent): # pylint: disable=W0622 684 """First array elements raised to second array elements :math:`{base}^{exponent}`, element-wise. 685 686 The result type of two scalar operands is determined as follows: 687 - If the exponent is an integral value, then the result type takes the type of the base. 688 - Otherwise, the result type follows 689 [Implicit type casting in binary operations](https://docs.taichi-lang.org/docs/type#implicit-type-casting-in-binary-operations). 690 691 With the above rules, an integral value raised to a negative integral value cannot have a 692 feasible type. Therefore, an exception will be raised if debug mode or optimization passes 693 are on; otherwise 1 will be returned. 694 695 In the following situations, the result is undefined: 696 - A negative value raised to a non-integral value. 697 - A zero value raised to a non-positive value. 698 699 Args: 700 base (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 701 The bases. 702 exponent (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 703 The exponents. 704 705 Returns: 706 `base` raised to `exponent`. This is a scalar if both `base` and `exponent` are scalars. 707 708 Example:: 709 710 >>> @ti.kernel 711 >>> def test(): 712 >>> x = ti.Matrix([-2.0, 2.0]) 713 >>> y = -3 714 >>> z = ti.pow(x, y) 715 >>> print(z) 716 >>> 717 >>> test() 718 [-0.125000, 0.125000] 719 """ 720 return _binary_operation(_ti_core.expr_pow, _bt_ops_mod.pow, base, exponent)
First array elements raised to second array elements \( {base}^{exponent} \), element-wise.
The result type of two scalar operands is determined as follows:
- If the exponent is an integral value, then the result type takes the type of the base.
- Otherwise, the result type follows Implicit type casting in binary operations.
With the above rules, an integral value raised to a negative integral value cannot have a feasible type. Therefore, an exception will be raised if debug mode or optimization passes are on; otherwise 1 will be returned.
In the following situations, the result is undefined:
- A negative value raised to a non-integral value.
- A zero value raised to a non-positive value.
Args:
base (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The bases.
exponent (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The exponents.
Returns:
base raised to exponent. This is a scalar if both base and exponent are scalars.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([-2.0, 2.0])
>>> y = -3
>>> z = ti.pow(x, y)
>>> print(z)
>>>
>>> test()
[-0.125000, 0.125000]
228def pyfunc(fn: Callable) -> TaichiCallable: 229 """Marks a function as callable in both Taichi and Python scopes. 230 231 When called inside the Taichi scope, Taichi will JIT compile it into 232 native instructions. Otherwise it will be invoked directly as a 233 Python function. 234 235 See also :func:`~taichi.lang.kernel_impl.func`. 236 237 Args: 238 fn (Callable): The Python function to be decorated 239 240 Returns: 241 Callable: The decorated function 242 """ 243 is_classfunc = _inside_class(level_of_class_stackframe=3) 244 fun = Func(fn, _classfunc=is_classfunc, _pyfunc=True) 245 taichi_callable = TaichiCallable(fn, fun) 246 taichi_callable._is_taichi_function = True 247 taichi_callable._is_real_function = False 248 return taichi_callable
Marks a function as callable in both Taichi and Python scopes.
When called inside the Taichi scope, Taichi will JIT compile it into native instructions. Otherwise it will be invoked directly as a Python function.
See also ~taichi.lang.kernel_impl.func().
Args: fn (Callable): The Python function to be decorated
Returns: Callable: The decorated function
563def random(dtype=float) -> Union[float, int]: 564 """Return a single random float/integer according to the specified data type. 565 Must be called in taichi scope. 566 567 If the required `dtype` is float type, this function returns a random number 568 sampled from the uniform distribution in the half-open interval [0, 1). 569 570 For integer types this function returns a random integer in the 571 half-open interval [0, 2^32) if a 32-bit integer is required, 572 or a random integer in the half-open interval [0, 2^64) if a 573 64-bit integer is required. 574 575 Args: 576 dtype (:mod:`~taichi.types.primitive_types`): Type of the required random value. 577 578 Returns: 579 A random value with type `dtype`. 580 581 Example:: 582 583 >>> @ti.kernel 584 >>> def test(): 585 >>> x = ti.random(float) 586 >>> print(x) # 0.090257 587 >>> 588 >>> y = ti.random(ti.f64) 589 >>> print(y) # 0.716101627301 590 >>> 591 >>> i = ti.random(ti.i32) 592 >>> print(i) # -963722261 593 >>> 594 >>> j = ti.random(ti.i64) 595 >>> print(j) # 73412986184350777 596 """ 597 dtype = cook_dtype(dtype) 598 x = expr.Expr(_ti_core.make_rand_expr(dtype, _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()))) 599 return impl.expr_init(x)
Return a single random float/integer according to the specified data type. Must be called in taichi scope.
If the required dtype is float type, this function returns a random number
sampled from the uniform distribution in the half-open interval [0, 1).
For integer types this function returns a random integer in the half-open interval [0, 2^32) if a 32-bit integer is required, or a random integer in the half-open interval [0, 2^64) if a 64-bit integer is required.
Args:
dtype (~taichi.types.primitive_types): Type of the required random value.
Returns:
A random value with type dtype.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.random(float)
>>> print(x) # 0.090257
>>>
>>> y = ti.random(ti.f64)
>>> print(y) # 0.716101627301
>>>
>>> i = ti.random(ti.i32)
>>> print(i) # -963722261
>>>
>>> j = ti.random(ti.i64)
>>> print(j) # 73412986184350777
804def raw_div(x1, x2): 805 """Return `x1 // x2` if both `x1`, `x2` are integers, otherwise return `x1/x2`. 806 807 Args: 808 x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): Dividend. 809 x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): Divisor. 810 811 Returns: 812 Return `x1 // x2` if both `x1`, `x2` are integers, otherwise return `x1/x2`. 813 814 Example:: 815 816 >>> @ti.kernel 817 >>> def main(): 818 >>> x = 5 819 >>> y = 3 820 >>> print(raw_div(x, y)) # 1 821 >>> z = 4.0 822 >>> print(raw_div(x, z)) # 1.25 823 """ 824 825 def c_div(a, b): 826 if isinstance(a, int) and isinstance(b, int): 827 return a // b 828 return a / b 829 830 return _binary_operation(_ti_core.expr_div, c_div, x1, x2)
Return x1 // x2 if both x1, x2 are integers, otherwise return x1/x2.
Args:
x1 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Dividend.
x2 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Divisor.
Returns:
Return x1 // x2 if both x1, x2 are integers, otherwise return x1/x2.
Example::
>>> @ti.kernel
>>> def main():
>>> x = 5
>>> y = 3
>>> print(raw_div(x, y)) # 1
>>> z = 4.0
>>> print(raw_div(x, z)) # 1.25
833def raw_mod(x1, x2): 834 """Return the remainder of `x1/x2`, element-wise. 835 This is the C-style `mod` function. 836 837 Args: 838 x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 839 The dividend. 840 x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 841 The divisor. 842 843 Returns: 844 The remainder of `x1` divided by `x2`. 845 846 Example:: 847 848 >>> @ti.kernel 849 >>> def main(): 850 >>> print(ti.mod(-4, 3)) # 2 851 >>> print(ti.raw_mod(-4, 3)) # -1 852 """ 853 854 def c_mod(x, y): 855 return x - y * int(float(x) / y) 856 857 return _binary_operation(_ti_core.expr_mod, c_mod, x1, x2)
Return the remainder of x1/x2, element-wise.
This is the C-style mod function.
Args:
x1 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The dividend.
x2 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The divisor.
Returns:
The remainder of x1 divided by x2.
Example::
>>> @ti.kernel
>>> def main():
>>> print(ti.mod(-4, 3)) # 2
>>> print(ti.raw_mod(-4, 3)) # -1
339def rescale_index(a, b, I): 340 """Rescales the index 'I' of field (or SNode) 'a' to match the shape of SNode 'b'. 341 342 Args: 343 344 a, b (Union[:class:`~taichi.Field`, :class:`~taichi.MatrixField`): Input taichi fields or snodes. 345 I (Union[list, :class:`~taichi.Vector`]): grouped loop index. 346 347 Returns: 348 Ib (:class:`~taichi.Vector`): rescaled grouped loop index 349 """ 350 351 assert isinstance(a, (Field, SNode)), "The first argument must be a field or an SNode" 352 assert isinstance(b, (Field, SNode)), "The second argument must be a field or an SNode" 353 if isinstance(I, list): 354 n = len(I) 355 else: 356 assert isinstance( 357 I, (expr.Expr, matrix.Matrix) 358 ), "The third argument must be an index (list, ti.Vector, or Expr with TensorType)" 359 n = I.n 360 361 from taichi.lang.kernel_impl import pyfunc # pylint: disable=C0415 362 363 @pyfunc 364 def _rescale_index(): 365 result = matrix.Vector([I[i] for i in range(n)]) 366 for i in impl.static(range(min(n, min(len(a.shape), len(b.shape))))): 367 if a.shape[i] > b.shape[i]: 368 result[i] = I[i] // (a.shape[i] // b.shape[i]) 369 if a.shape[i] < b.shape[i]: 370 result[i] = I[i] * (b.shape[i] // a.shape[i]) 371 return result 372 373 return _rescale_index()
Rescales the index 'I' of field (or SNode) 'a' to match the shape of SNode 'b'.
Args:
a, b (Union[`~taichi.Field`, `~taichi.MatrixField`): Input taichi fields or snodes.
I (Union[list, `~taichi.Vector`]): grouped loop index.
Returns:
Ib (~taichi.Vector): rescaled grouped loop index
206def reset(): 207 """Resets Taichi to its initial state. 208 This will destroy all the allocated fields and kernels, and restore 209 the runtime to its default configuration. 210 211 Example:: 212 213 >>> a = ti.field(ti.i32, shape=()) 214 >>> a[None] = 1 215 >>> print("before reset: ", a) 216 before rest: 1 217 >>> 218 >>> ti.reset() 219 >>> print("after reset: ", a) 220 # will raise error because a is unavailable after reset. 221 """ 222 impl.reset() 223 global runtime 224 runtime = impl.get_runtime()
Resets Taichi to its initial state. This will destroy all the allocated fields and kernels, and restore the runtime to its default configuration.
Example::
>>> a = ti.field(ti.i32, shape=())
>>> a[None] = 1
>>> print("before reset: ", a)
before rest: 1
>>>
>>> ti.reset()
>>> print("after reset: ", a)
# will raise error because a is unavailable after reset.
311def round(x, dtype=None): # pylint: disable=redefined-builtin 312 """Round to the nearest integer, element-wise. 313 314 Args: 315 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 316 A scalar or a matrix. 317 318 dtype: (:mod:`~taichi.types.primitive_types`): the returned type, default to `None`. If \ 319 set to `None` the retuned value will have the same type with `x`. 320 321 Returns: 322 The nearest integer of `x`, with return value type `dtype`. 323 324 Example:: 325 326 >>> @ti.kernel 327 >>> def test(): 328 >>> x = ti.Vector([-1.5, 1.2, 2.7]) 329 >>> print(ti.round(x)) 330 [-2., 1., 3.] 331 """ 332 result = _round(x) 333 if dtype is not None: 334 result = cast(result, dtype) 335 return result
Round to the nearest integer, element-wise.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): A scalar or a matrix.
dtype: (`~taichi.types.primitive_types`): the returned type, default to `None`. If set to `None` the retuned value will have the same type with `x`.
Returns:
The nearest integer of x, with return value type dtype.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([-1.5, 1.2, 2.7])
>>> print(ti.round(x))
[-2., 1., 3.]
290def rsqrt(x): 291 """The reciprocal of the square root function. 292 293 Args: 294 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 295 A scalar or a matrix. 296 297 Returns: 298 The reciprocal of `sqrt(x)`. 299 """ 300 301 def _rsqrt(x): 302 return 1 / np.sqrt(x) 303 304 return _unary_operation(_ti_core.expr_rsqrt, _rsqrt, x)
The reciprocal of the square root function.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): A scalar or a matrix.
Returns:
The reciprocal of sqrt(x).
1070def select(cond, x1, x2): 1071 """Return an array drawn from elements in `x1` or `x2`, 1072 depending on the conditions in `cond`. 1073 1074 Args: 1075 cond (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1076 The array of conditions. 1077 x1, x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 1078 The arrays where the output elements are taken from. 1079 1080 Returns: 1081 The output at position `k` is the k-th element of `x1` if the k-th element 1082 in `cond` is `True`, otherwise it's the k-th element of `x2`. 1083 1084 Example:: 1085 1086 >>> @ti.kernel 1087 >>> def main(): 1088 >>> cond = ti.Matrix([0, 1, 0, 1]) 1089 >>> x = ti.Matrix([1, 2, 3, 4]) 1090 >>> y = ti.Matrix([-1, -2, -3, -4]) 1091 >>> print(ti.select(cond, x, y)) 1092 >>> 1093 >>> main() 1094 [-1, 2, -3, 4] 1095 """ 1096 # TODO: systematically resolve `-1 = True` problem by introducing u1: 1097 cond = logical_not(logical_not(cond)) 1098 1099 def py_select(cond, x1, x2): 1100 return x1 * cond + x2 * (1 - cond) 1101 1102 return _ternary_operation(_ti_core.expr_select, py_select, cond, x1, x2)
Return an array drawn from elements in x1 or x2,
depending on the conditions in cond.
Args:
cond (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The array of conditions.
x1, x2 (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The arrays where the output elements are taken from.
Returns:
The output at position k is the k-th element of x1 if the k-th element
in cond is True, otherwise it's the k-th element of x2.
Example::
>>> @ti.kernel
>>> def main():
>>> cond = ti.Matrix([0, 1, 0, 1])
>>> x = ti.Matrix([1, 2, 3, 4])
>>> y = ti.Matrix([-1, -2, -3, -4])
>>> print(ti.select(cond, x, y))
>>>
>>> main()
[-1, 2, -3, 4]
179def sin(x): 180 """Trigonometric sine, element-wise. 181 182 Args: 183 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 184 Angle, in radians. 185 186 Returns: 187 The sine of each element of `x`. 188 189 Example:: 190 191 >>> from math import pi 192 >>> x = ti.Matrix([-pi/2., 0, pi/2.]) 193 >>> ti.sin(x) 194 [-1., 0., 1.] 195 """ 196 return _unary_operation(_ti_core.expr_sin, np.sin, x)
Trigonometric sine, element-wise.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Angle, in radians.
Returns:
The sine of each element of x.
Example::
>>> from math import pi
>>> x = ti.Matrix([-pi/2., 0, pi/2.])
>>> ti.sin(x)
[-1., 0., 1.]
269def sqrt(x): 270 """Return the non-negative square-root of a scalar or a matrix, 271 element wise. If `x < 0` an exception is raised. 272 273 Args: 274 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 275 The scalar or matrix whose square-roots are required. 276 277 Returns: 278 The square-root `y` so that `y >= 0` and `y^2 = x`. `y` has the same type as `x`. 279 280 Example:: 281 282 >>> x = ti.Matrix([1., 4., 9.]) 283 >>> y = ti.sqrt(x) 284 >>> y 285 [1.0, 2.0, 3.0] 286 """ 287 return _unary_operation(_ti_core.expr_sqrt, np.sqrt, x)
Return the non-negative square-root of a scalar or a matrix,
element wise. If x < 0 an exception is raised.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The scalar or matrix whose square-roots are required.
Returns:
The square-root y so that y >= 0 and y^2 = x. y has the same type as x.
Example::
>>> x = ti.Matrix([1., 4., 9.])
>>> y = ti.sqrt(x)
>>> y
[1.0, 2.0, 3.0]
1091def static(x, *xs) -> Any: 1092 """Evaluates a Taichi-scope expression at compile time. 1093 1094 `static()` is what enables the so-called metaprogramming in Taichi. It is 1095 in many ways similar to ``constexpr`` in C++. 1096 1097 See also https://docs.taichi-lang.org/docs/meta. 1098 1099 Args: 1100 x (Any): an expression to be evaluated 1101 *xs (Any): for Python-ish swapping assignment 1102 1103 Example: 1104 The most common usage of `static()` is for compile-time evaluation:: 1105 1106 >>> cond = False 1107 >>> 1108 >>> @ti.kernel 1109 >>> def run(): 1110 >>> if ti.static(cond): 1111 >>> do_a() 1112 >>> else: 1113 >>> do_b() 1114 1115 Depending on the value of ``cond``, ``run()`` will be directly compiled 1116 into either ``do_a()`` or ``do_b()``. Thus there won't be a runtime 1117 condition check. 1118 1119 Another common usage is for compile-time loop unrolling:: 1120 1121 >>> @ti.kernel 1122 >>> def run(): 1123 >>> for i in ti.static(range(3)): 1124 >>> print(i) 1125 >>> 1126 >>> # The above will be unrolled to: 1127 >>> @ti.kernel 1128 >>> def run(): 1129 >>> print(0) 1130 >>> print(1) 1131 >>> print(2) 1132 """ 1133 if len(xs): # for python-ish pointer assign: x, y = ti.static(y, x) 1134 return [static(x)] + [static(x) for x in xs] 1135 1136 if ( 1137 isinstance( 1138 x, 1139 ( 1140 bool, 1141 int, 1142 float, 1143 range, 1144 list, 1145 tuple, 1146 enumerate, 1147 GroupedNDRange, 1148 _Ndrange, 1149 zip, 1150 filter, 1151 map, 1152 ), 1153 ) 1154 or x is None 1155 ): 1156 return x 1157 if isinstance(x, (np.bool_, np.integer, np.floating)): 1158 return x 1159 1160 if isinstance(x, AnyArray): 1161 return x 1162 if isinstance(x, Field): 1163 return x 1164 if isinstance(x, (FunctionType, MethodType, BoundTaichiCallable, TaichiCallable)): 1165 return x 1166 raise ValueError(f"Input to ti.static must be compile-time constants or global pointers, instead of {type(x)}")
Evaluates a Taichi-scope expression at compile time.
static() is what enables the so-called metaprogramming in Taichi. It is
in many ways similar to constexpr in C++.
See also https://docs.taichi-lang.org/docs/meta.
Args: x (Any): an expression to be evaluated *xs (Any): for Python-ish swapping assignment
Example:
The most common usage of static() is for compile-time evaluation::
>>> cond = False
>>>
>>> @ti.kernel
>>> def run():
>>> if ti.static(cond):
>>> do_a()
>>> else:
>>> do_b()
Depending on the value of ``cond``, ``run()`` will be directly compiled
into either ``do_a()`` or ``do_b()``. Thus there won't be a runtime
condition check.
Another common usage is for compile-time loop unrolling::
>>> @ti.kernel
>>> def run():
>>> for i in ti.static(range(3)):
>>> print(i)
>>>
>>> # The above will be unrolled to:
>>> @ti.kernel
>>> def run():
>>> print(0)
>>> print(1)
>>> print(2)
560def static_assert(cond, msg=None): 561 """Throw AssertionError when `cond` is False. 562 563 This function is called at compile time and has no runtime overhead. 564 The bool value in `cond` must can be determined at compile time. 565 566 Args: 567 cond (bool): an expression with a bool value. 568 msg (str): assertion message. 569 570 Example:: 571 572 >>> year = 2001 573 >>> @ti.kernel 574 >>> def test(): 575 >>> ti.static_assert(year % 4 == 0, "the year must be a lunar year") 576 AssertionError: the year must be a lunar year 577 """ 578 if isinstance(cond, Expr): 579 raise TaichiTypeError("Static assert with non-static condition") 580 if msg is not None: 581 assert cond, msg 582 else: 583 assert cond
Throw AssertionError when cond is False.
This function is called at compile time and has no runtime overhead.
The bool value in cond must can be determined at compile time.
Args: cond (bool): an expression with a bool value. msg (str): assertion message.
Example::
>>> year = 2001
>>> @ti.kernel
>>> def test():
>>> ti.static_assert(year % 4 == 0, "the year must be a lunar year")
AssertionError: the year must be a lunar year
550@taichi_scope 551def static_print(*args, __p=print, **kwargs): 552 """The print function in Taichi scope. 553 554 This function is called at compile time and has no runtime overhead. 555 """ 556 __p(*args, **kwargs)
The print function in Taichi scope.
This function is called at compile time and has no runtime overhead.
1195def stop_grad(x): 1196 """Stops computing gradients during back propagation. 1197 1198 Args: 1199 x (:class:`~taichi.Field`): A field. 1200 """ 1201 compiling_callable = get_runtime().compiling_callable 1202 assert compiling_callable is not None 1203 compiling_callable.ast_builder().stop_grad(x.snode.ptr)
Stops computing gradients during back propagation.
Args:
x (~taichi.Field): A field.
7def sync(): 8 """Blocks the calling thread until all the previously 9 launched Taichi kernels have completed. 10 """ 11 impl.get_runtime().sync()
Blocks the calling thread until all the previously launched Taichi kernels have completed.
406def tan(x): 407 """Trigonometric tangent function, element-wise. 408 409 Equivalent to `ti.sin(x)/ti.cos(x)` element-wise. 410 411 Args: 412 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 413 Input scalar or matrix. 414 415 Returns: 416 The tangent values of `x`. 417 418 Example:: 419 420 >>> from math import pi 421 >>> @ti.kernel 422 >>> def test(): 423 >>> x = ti.Matrix([-pi, pi/2, pi]) 424 >>> y = ti.tan(x) 425 >>> print(y) 426 >>> 427 >>> test() 428 [-0.0, -22877334.0, 0.0] 429 """ 430 return _unary_operation(_ti_core.expr_tan, np.tan, x)
Trigonometric tangent function, element-wise.
Equivalent to ti.sin(x)/ti.cos(x) element-wise.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
Returns:
The tangent values of x.
Example::
>>> from math import pi
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([-pi, pi/2, pi])
>>> y = ti.tan(x)
>>> print(y)
>>>
>>> test()
[-0.0, -22877334.0, 0.0]
433def tanh(x): 434 """Compute the hyperbolic tangent of `x`, element-wise. 435 436 Args: 437 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ 438 Input scalar or matrix. 439 440 Returns: 441 The corresponding hyperbolic tangent values. 442 443 Example:: 444 445 >>> @ti.kernel 446 >>> def test(): 447 >>> x = ti.Matrix([-1.0, 0.0, 1.0]) 448 >>> y = ti.tanh(x) 449 >>> print(y) 450 >>> 451 >>> test() 452 [-0.761594, 0.000000, 0.761594] 453 """ 454 return _unary_operation(_ti_core.expr_tanh, np.tanh, x)
Compute the hyperbolic tangent of x, element-wise.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): Input scalar or matrix.
Returns: The corresponding hyperbolic tangent values.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([-1.0, 0.0, 1.0])
>>> y = ti.tanh(x)
>>> print(y)
>>>
>>> test()
[-0.761594, 0.000000, 0.761594]
1029@taichi_scope 1030def zero(x): 1031 # TODO: get dtype from Expr and Matrix: 1032 """Returns an array of zeros with the same shape and type as the input. It's also a scalar 1033 if the input is a scalar. 1034 1035 Args: 1036 x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): The input. 1037 1038 Returns: 1039 A new copy of the input but filled with zeros. 1040 1041 Example:: 1042 1043 >>> x = ti.Vector([1, 1]) 1044 >>> @ti.kernel 1045 >>> def test(): 1046 >>> y = ti.zero(x) 1047 >>> print(y) 1048 [0, 0] 1049 """ 1050 return x * 0
Returns an array of zeros with the same shape and type as the input. It's also a scalar if the input is a scalar.
Args:
x (Union[~taichi.types.primitive_types, ~taichi.Matrix]): The input.
Returns: A new copy of the input but filled with zeros.
Example::
>>> x = ti.Vector([1, 1])
>>> @ti.kernel
>>> def test():
>>> y = ti.zero(x)
>>> print(y)
[0, 0]